mirror of
https://github.com/qdrant/fastembed.git
synced 2026-09-23 14:37:51 -05:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Nirant Kasliwal
|
2b485e6e6a | ||
|
Nirant Kasliwal
|
b5bce4c2a1 |
@@ -1,57 +0,0 @@
|
||||
name: Bug/New Model Request
|
||||
description: File a bug report/Request a new Model
|
||||
title: "[Bug/Model Request]: "
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thanks for taking the time to fill out this bug report!
|
||||
- type: textarea
|
||||
id: what-happened
|
||||
attributes:
|
||||
label: What happened?
|
||||
description: Also tell us, what did you expect to happen?
|
||||
placeholder: Tell us what you see!
|
||||
value: "A bug happened!"
|
||||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
id: python-version
|
||||
attributes:
|
||||
label: What Python version are you on? e.g. python --version
|
||||
description: Also tell us, what package manager are you using e.g. conda, pip, poetry?
|
||||
placeholder: Python3.10
|
||||
validations:
|
||||
required: true
|
||||
- type: dropdown
|
||||
id: version
|
||||
attributes:
|
||||
label: Version
|
||||
description: What version of FastEmbed are you running? python -c "import fastembed; print(fastembed.__version__)". If you're not on the latest, please upgrade and see if the problem persists.
|
||||
options:
|
||||
- 0.2.7 (Latest)
|
||||
- 0.2.6
|
||||
- 0.2.5
|
||||
- 0.2.4
|
||||
- 0.2.3
|
||||
- 0.2.2
|
||||
- 0.2.1
|
||||
- 0.1.x
|
||||
default: 0
|
||||
validations:
|
||||
required: true
|
||||
- type: dropdown
|
||||
id: os
|
||||
attributes:
|
||||
label: What os are you seeing the problem on?
|
||||
multiple: true
|
||||
options:
|
||||
- Linux
|
||||
- MacOS
|
||||
- Windows
|
||||
- type: textarea
|
||||
id: logs
|
||||
attributes:
|
||||
label: Relevant stack traces and/or logs
|
||||
description: Please copy and paste any relevant raised exceptions. This will be automatically formatted into code, so no need for backticks.
|
||||
render: shell
|
||||
@@ -1,5 +0,0 @@
|
||||
blank_issues_enabled: false
|
||||
contact_links:
|
||||
- name: GitHub Community Support
|
||||
url: https://github.com/qdrant/fastembed/discussions
|
||||
about: Please ask and answer questions here.
|
||||
@@ -15,6 +15,7 @@ on:
|
||||
tags:
|
||||
- 'v*' # Push events to every version tag
|
||||
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
|
||||
|
||||
@@ -1,11 +1,8 @@
|
||||
name: Tests
|
||||
run-name: Tests (gpu)
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ master, main, gpu ]
|
||||
schedule:
|
||||
- cron: 0 0 * * *
|
||||
branches: [ master, main ]
|
||||
pull_request:
|
||||
|
||||
env:
|
||||
@@ -24,30 +21,26 @@ jobs:
|
||||
- '3.12.x'
|
||||
os:
|
||||
- ubuntu-latest
|
||||
- macos-latest
|
||||
- windows-latest
|
||||
|
||||
runs-on: ${{ matrix.os }}
|
||||
|
||||
name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
- uses: actions/checkout@v2
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
uses: actions/setup-python@v2
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install poetry
|
||||
poetry config virtualenvs.create false
|
||||
poetry install --no-interaction --no-ansi --without docs
|
||||
|
||||
- name: Install Test Dependencies
|
||||
run: pip install pytest pytest-md pytest-emoji
|
||||
|
||||
- name: Run pytest
|
||||
uses: pavelzw/pytest-action@v2
|
||||
with:
|
||||
verbose: true
|
||||
emoji: true
|
||||
job-summary: true
|
||||
report-title: 'FastEmbed Test Report'
|
||||
poetry install --no-interaction --no-ansi
|
||||
- name: Run tests
|
||||
run: |
|
||||
export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
|
||||
pytest
|
||||
shell: bash
|
||||
|
||||
+40
-4
@@ -85,8 +85,28 @@ ipython_config.py
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
.python-version
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/#use-with-ide
|
||||
.pdm.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
@@ -132,11 +152,27 @@ dmypy.json
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
.idea/
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
.DS_Store
|
||||
nbs/*.tar.gz
|
||||
*.tar.gz
|
||||
**/local_cache/
|
||||
nbs/fast-*/*
|
||||
local_cache/*/*
|
||||
*/local_cache/*/*
|
||||
*/*/local_cache/*/*
|
||||
docs/experimental/*.parquet
|
||||
docs/experimental/*.bin
|
||||
qdrant_storage/*
|
||||
experiments/models/*
|
||||
fooling_around/fast-multilingual-e5-large/config.json
|
||||
fooling_around/fast-multilingual-e5-large/model_optimized.onnx
|
||||
fooling_around/fast-multilingual-e5-large/model_optimized.onnx.data
|
||||
fooling_around/fast-multilingual-e5-large/ort_config.json
|
||||
fooling_around/fast-multilingual-e5-large/sentencepiece.bpe.model
|
||||
fooling_around/fast-multilingual-e5-large/special_tokens_map.json
|
||||
fooling_around/fast-multilingual-e5-large/tokenizer_config.json
|
||||
fooling_around/fast-multilingual-e5-large/tokenizer.json
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.3.4
|
||||
rev: v0.1.13
|
||||
hooks:
|
||||
- id: ruff
|
||||
types_or: [ python, pyi, jupyter ]
|
||||
|
||||
@@ -1,78 +0,0 @@
|
||||
# Contributing to FastEmbed!
|
||||
|
||||
:+1::tada: First off, thanks for taking the time to contribute! :tada::+1:
|
||||
|
||||
The following is a set of guidelines for contributing to FastEmbed. These are mostly guidelines, not rules. Use your best judgment, and feel free to propose changes to this document in a pull request.
|
||||
|
||||
## Table Of Contents
|
||||
|
||||
[I don't want to read this whole thing, I just have a question!!!](#i-dont-want-to-read-this-whole-thing-i-just-have-a-question)
|
||||
|
||||
[How Can I Contribute?](#how-can-i-contribute)
|
||||
* [Your First Code Contribution](#your-first-code-contribution)
|
||||
* [Adding New Models](#adding-new-models)
|
||||
|
||||
[Styleguides](#styleguides)
|
||||
* [Code Lint](#code-lint)
|
||||
* [Pre-Commit Hooks](#pre-commit-hooks)
|
||||
|
||||
## I don't want to read this whole thing I just have a question!!!
|
||||
|
||||
> **Note:** Please don't file an issue to ask a question. You'll get faster results by using the resources below:
|
||||
|
||||
* [FastEmbed Docs](https://qdrant.github.io/fastembed/)
|
||||
* [Qdrant Discord](https://discord.gg/Qy6HCJK9Dc)
|
||||
|
||||
## How Can I Contribute?
|
||||
|
||||
## How Do I Submit A (Good) Bug Report?
|
||||
|
||||
Bugs are tracked as [GitHub issues](https://guides.github.com/features/issues/).
|
||||
|
||||
Explain the problem and include additional details to help maintainers reproduce the problem:
|
||||
|
||||
* **Use a clear and descriptive title** for the issue to identify the problem.
|
||||
* **Describe the exact steps which reproduce the problem** in as many details as possible. For example, start by explaining how you are using FastEmbed, e.g. with Langchain, Qdrant Client, Llama Index and which command exactly you used. When listing steps, **don't just say what you did, but explain how you did it**.
|
||||
* **Provide specific examples to demonstrate the steps**. Include links to files or GitHub projects, or copy/pasteable snippets, which you use in those examples. If you're providing snippets in the issue, use [Markdown code blocks](https://help.github.com/articles/markdown-basics/#multiple-lines).
|
||||
* **Describe the behavior you observed after following the steps** and point out what exactly is the problem with that behavior.
|
||||
* **Explain which behavior you expected to see instead and why.**
|
||||
* **If the problem is related to performance or memory**, include a [call stack profile capture](https://github.com/joerick/pyinstrument) and your observations.
|
||||
|
||||
Include details about your configuration and environment:
|
||||
|
||||
* **Which version of FastEmbed are you using?** You can get the exact version by running `python -c "import fastembed; print(fastembed.__version__)"`.
|
||||
* **What's the name and version of the OS you're using**?
|
||||
* **Which packages do you have installed?** You can get that list by running `pip freeze`
|
||||
|
||||
### Your First Code Contribution
|
||||
|
||||
Unsure where to begin contributing to FastEmbed? You can start by looking through these `good-first-issue`issues:
|
||||
|
||||
* [Good First Issue](https://github.com/qdrant/fastembed/labels/good%20first%20issue) - issues which should only require a few lines of code, and a test or two. These are a great way to get started with FastEmbed. This includes adding new models which are already tested and ready on Huggingface Hub.
|
||||
|
||||
## Pull Requests
|
||||
|
||||
The best way to learn about the mechanics of FastEmbed is to start working on it.
|
||||
|
||||
### Your First Code Contribution
|
||||
Your first code contribution can be small bug fixes:
|
||||
1. This PR adds a small bug fix for a single input: https://github.com/qdrant/fastembed/pull/148
|
||||
2. This PR adds a check for the right file location and extension, specific to an OS: https://github.com/qdrant/fastembed/pull/128
|
||||
|
||||
Even documentation improvements and tests are most welcome:
|
||||
1. This PR fixes a README link: https://github.com/qdrant/fastembed/pull/143
|
||||
|
||||
### Adding New Models
|
||||
1. Open Requests for New Models are [here](https://github.com/qdrant/fastembed/labels/model%20request).
|
||||
2. There are quite a few pull requests that were merged for this purpose and you can use them as a reference. Here is an example: https://github.com/qdrant/fastembed/pull/129
|
||||
3. Make sure to add tests for the new model
|
||||
- The CANONICAL_VECTOR values must come from a reference implementation usually from Huggingface Transformers or Sentence Transformers
|
||||
- Here is a reference [Colab Notebook](https://colab.research.google.com/drive/1tNdV3DsiwsJzu2AXnUnoeF5av1Hp8HF1?usp=sharing) for how we will evaluate whether your VECTOR values in the test are correct or not.
|
||||
|
||||
## Styleguides
|
||||
|
||||
### Code Lint
|
||||
We use ruff for code linting. It should be installed with poetry since it's a dev dependency.
|
||||
|
||||
### Pre-Commit Hooks
|
||||
We use pre-commit hooks to ensure that the code is linted before it's committed. You can install pre-commit hooks by running `pre-commit install` in the root directory of the project.
|
||||
@@ -2,72 +2,42 @@
|
||||
|
||||
FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
|
||||
|
||||
The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/).
|
||||
The default text embedding (`TextEmbedding`) model is Flag Embedding, the top model in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
|
||||
|
||||
## 📈 Why FastEmbed?
|
||||
1. Light & Fast
|
||||
- Quantized model weights
|
||||
- ONNX Runtime, no PyTorch dependency
|
||||
- CPU-first design
|
||||
- Data-parallelism for encoding of large datasets
|
||||
|
||||
1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda.
|
||||
|
||||
2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
|
||||
|
||||
3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [supported](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
|
||||
2. Accuracy/Recall
|
||||
- Better than OpenAI Ada-002
|
||||
- Default is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard
|
||||
- List of [supported models](https://qdrant.github.io/fastembed/examples/Supported_Models/) - including multilingual models
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
To install the FastEmbed library, pip works best. You can install it with or without GPU support:
|
||||
To install the FastEmbed library, pip works:
|
||||
|
||||
```bash
|
||||
pip install fastembed
|
||||
```
|
||||
|
||||
### ⚡️ With GPU
|
||||
|
||||
```bash
|
||||
pip install fastembed-gpu
|
||||
```
|
||||
|
||||
## 📖 Quickstart
|
||||
|
||||
```python
|
||||
from fastembed import TextEmbedding
|
||||
from typing import List
|
||||
import numpy as np
|
||||
|
||||
# Example list of documents
|
||||
documents: List[str] = [
|
||||
"This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
|
||||
"fastembed is supported by and maintained by Qdrant.",
|
||||
"passage: Hello, World!",
|
||||
"query: Hello, World!", # these are two different embedding
|
||||
"passage: This is an example passage.",
|
||||
"fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
|
||||
]
|
||||
|
||||
# This will trigger the model download and initialization
|
||||
embedding_model = TextEmbedding()
|
||||
print("The model BAAI/bge-small-en-v1.5 is ready to use.")
|
||||
|
||||
embeddings_generator = embedding_model.embed(documents) # reminder this is a generator
|
||||
embeddings_list = list(embedding_model.embed(documents))
|
||||
# you can also convert the generator to a list, and that to a numpy array
|
||||
len(embeddings_list[0]) # Vector of 384 dimensions
|
||||
```
|
||||
|
||||
### ⚡️ FastEmbed on a GPU
|
||||
|
||||
FastEmbed supports running on GPU devices.
|
||||
It requires installation of the `fastembed-gpu` package.
|
||||
|
||||
```bash
|
||||
pip install fastembed-gpu
|
||||
```
|
||||
|
||||
Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
|
||||
|
||||
```python
|
||||
from fastembed import TextEmbedding
|
||||
|
||||
embedding_model = TextEmbedding(
|
||||
model_name="BAAI/bge-small-en-v1.5",
|
||||
providers=["CUDAExecutionProvider"]
|
||||
)
|
||||
print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
|
||||
|
||||
embedding_model = TextEmbedding(model_name="BAAI/bge-base-en")
|
||||
embeddings: List[np.ndarray] = list(embedding_model.embed(documents)) # Note the list() call - this is a generator
|
||||
```
|
||||
|
||||
## Usage with Qdrant
|
||||
@@ -78,13 +48,7 @@ Installation with Qdrant Client in Python:
|
||||
pip install qdrant-client[fastembed]
|
||||
```
|
||||
|
||||
or
|
||||
|
||||
```bash
|
||||
pip install qdrant-client[fastembed-gpu]
|
||||
```
|
||||
|
||||
You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
|
||||
You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
|
||||
|
||||
```python
|
||||
from qdrant_client import QdrantClient
|
||||
@@ -119,4 +83,8 @@ search_result = client.query(
|
||||
query_text="This is a query document"
|
||||
)
|
||||
print(search_result)
|
||||
```
|
||||
```
|
||||
|
||||
#### Similar Work
|
||||
|
||||
Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
|
||||
|
||||
-41
@@ -1,41 +0,0 @@
|
||||
# Releasing FastEmbed
|
||||
|
||||
This is a guide how to release `fastembed` and `fastembed-gpu` packages.
|
||||
|
||||
## How to
|
||||
|
||||
1. Accumulate changes in the `main` branch.
|
||||
2. Bump the version in `pyproject.toml`
|
||||
|
||||
3. Rebase the `gpu` branch on `main` and resolve conflicts if occurred:
|
||||
|
||||
```bash
|
||||
git checkout gpu
|
||||
git rebase main
|
||||
git push origin gpu
|
||||
```
|
||||
|
||||
4. Draft release notes
|
||||
5. Checkout to `main` and create a tag, e.g.:
|
||||
|
||||
```bash
|
||||
git checkout main
|
||||
git tag -a v0.1.0 -m "Release v0.1.0"
|
||||
```
|
||||
|
||||
6. Checkout `gpu` and create a tag, e.g.:
|
||||
|
||||
```bash
|
||||
git checkout gpu
|
||||
git tag -a v0.1.0-gpu -m "Release v0.1.0"
|
||||
```
|
||||
|
||||
7. Push tags:
|
||||
|
||||
```bash
|
||||
git push --tags
|
||||
```
|
||||
|
||||
8. Verify that both packages have been published successfully on PyPI. Try installing them and verify imports.
|
||||
9. Create a release on GitHub with the written release notes.
|
||||
|
||||
+138
-136
@@ -11,9 +11,7 @@
|
||||
"\n",
|
||||
"## Quick Start\n",
|
||||
"\n",
|
||||
"The fastembed package is designed to be easy to use. We'll be using `TextEmbedding` class. It takes a list of strings as input and returns a generator of vectors.\n",
|
||||
"\n",
|
||||
"> 💡 You can learn more about generators from [Python Wiki](https://wiki.python.org/moin/Generators)"
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -23,7 +21,15 @@
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install -Uqq fastembed"
|
||||
"!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:"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -33,115 +39,43 @@
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "890cc3b969354eec8d149d143e301a7a",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Fetching 9 files: 0%| | 0/9 [00:00<?, ?it/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
"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": [
|
||||
"The model BAAI/bge-small-en-v1.5 is ready to use.\n"
|
||||
"(384,)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"384"
|
||||
]
|
||||
},
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from typing import List\n",
|
||||
"\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"from fastembed import TextEmbedding\n",
|
||||
"\n",
|
||||
"from fastembed.embedding import DefaultEmbedding\n",
|
||||
"\n",
|
||||
"# Example list of documents\n",
|
||||
"documents: List[str] = [\n",
|
||||
" \"This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.\",\n",
|
||||
" \"Hello, World!\",\n",
|
||||
" \"This is an example document.\",\n",
|
||||
" \"fastembed is supported by and maintained by Qdrant.\",\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"# This will trigger the model download and initialization\n",
|
||||
"embedding_model = TextEmbedding()\n",
|
||||
"print(\"The model BAAI/bge-small-en-v1.5 is ready to use.\")\n",
|
||||
"\n",
|
||||
"embeddings_generator = embedding_model.embed(documents)\n",
|
||||
"embeddings_list = list(embeddings_generator)\n",
|
||||
"len(embeddings_list[0]) # Vector of 384 dimensions"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d772190b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"> 💡 **Why do we use generators?**\n",
|
||||
"> \n",
|
||||
"> We use them to save memory mostly. Instead of loading all the vectors into memory, we can load them one by one. This is useful when you have a large dataset and you don't want to load all the vectors at once."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "8a225cb8",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Document: This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.\n",
|
||||
"Vector of type: <class 'numpy.ndarray'> with shape: (384,)\n",
|
||||
"Document: fastembed is supported by and maintained by Qdrant.\n",
|
||||
"Vector of type: <class 'numpy.ndarray'> with shape: (384,)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings_generator = embedding_model.embed(documents)\n",
|
||||
"\n",
|
||||
"for doc, vector in zip(documents, embeddings_generator):\n",
|
||||
" print(\"Document:\", doc)\n",
|
||||
" print(f\"Vector of type: {type(vector)} with shape: {vector.shape}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "769a1be9",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"(2, 384)"
|
||||
]
|
||||
},
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings_list = np.array(list(embedding_model.embed(documents)))\n",
|
||||
"embeddings_list.shape"
|
||||
"# Initialize the DefaultEmbedding class\n",
|
||||
"embedding_model = DefaultEmbedding()\n",
|
||||
"embeddings: List[np.ndarray] = list(embedding_model.embed(documents))\n",
|
||||
"print(embeddings[0].shape)"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -149,74 +83,142 @@
|
||||
"id": "8c49ae50",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We're using [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) a state of the art Flag Embedding model. The model does better than OpenAI text-embedding-ada-002. We've made it even faster by converting it to ONNX format and quantizing the model for you.\n",
|
||||
"## 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",
|
||||
"\n",
|
||||
"1. List of Strings: Your documents must be in a list, and each document must be a string\n",
|
||||
"2. For Retrieval Tasks with our default: If you're working with queries and passages, you can add special labels to them:\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\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",
|
||||
"## Beyond the default model\n",
|
||||
"This will download the model weights, decompress to directory `local_cache` and load them into the Embedding class.\n",
|
||||
"\n",
|
||||
"The default model is built for speed and efficiency. If you need a more accurate model, you can use the `TextEmbedding` class to load any model from our list of available models. You can find the list of available models using `TextEmbedding.list_supported_models()`."
|
||||
"#### 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": "2e9c8766",
|
||||
"id": "8013eee9",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "9470ec542f3c4400a42452c2489a1abc",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Fetching 8 files: 0%| | 0/8 [00:00<?, ?it/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"100%|██████████| 4/4 [00:00<00:00, 361.82it/s]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"multilingual_large_model = TextEmbedding(\"intfloat/multilingual-e5-large\")"
|
||||
"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": "a9e70f0e",
|
||||
"id": "0d8c8e08",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"(4, 1024)"
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"(384,)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"np.array(\n",
|
||||
" list(multilingual_large_model.embed([\"Hello, world!\", \"你好世界\", \"¡Hola Mundo!\", \"नमस्ते!\"]))\n",
|
||||
").shape # Vector of 1024 dimensions"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "64fe20ed",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next: Checkout how to use FastEmbed with Qdrant for similarity search: [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/)"
|
||||
"print(embeddings[0].shape) # (384,) or similar output"
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -236,7 +238,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.13"
|
||||
"version": "3.9.17"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -1,410 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d14d29ebd3592ecb",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"# Late Interaction Text Embedding Models\n",
|
||||
"\n",
|
||||
"As of version 0.3.0 FastEmbed supports Late Interaction Text Embedding Models and currently available with one of the most popular embedding model of the family - ColBERT.\n",
|
||||
"\n",
|
||||
"## What is a Late Interaction Text Embedding Model?\n",
|
||||
"\n",
|
||||
"Late Interaction Text Embedding Model is a kind of information retrieval model which performs query and documents interactions at the scoring stage.\n",
|
||||
"In order to better understand it, we can compare it to the models without interaction. \n",
|
||||
"For instance, if you take a sentence-transformer model, compute embeddings for your documents, compute embeddings for your queries, and just compare them by cosine similarity, then you're retrieving points without interaction.\n",
|
||||
"\n",
|
||||
"It is a pretty much easy and straightforward approach, however we might be sacrificing some precision due to its simplicity. It is caused by several facts: \n",
|
||||
"- there is no interaction between queries and documents at the early stage (embedding generation) nor at the late stage (during scoring). \n",
|
||||
"- we are trying to encapsulate all the document information in only one pooled embedding, and obviously, some information might be lost.\n",
|
||||
"\n",
|
||||
"Late Interaction Text Embedding models are trying to address it by computing embeddings for each token in queries and documents, and then finding the most similar ones via model specific operation, e.g. ColBERT (Contextual Late Interaction over BERT) uses MaxSim operation.\n",
|
||||
"With this approach we can have not only a better representation of the documents, but also make queries and documents more aware one of another.\n",
|
||||
"\n",
|
||||
"For more information on ColBERT and MaxSim operation, you can check out [this blogpost](https://jina.ai/news/what-is-colbert-and-late-interaction-and-why-they-matter-in-search/) by Jina AI.\n",
|
||||
"\n",
|
||||
"## ColBERT in FastEmbed\n",
|
||||
"\n",
|
||||
"FastEmbed provides a simple way to use ColBERT model, similar to the ones it has with `TextEmbedding`.\n",
|
||||
" "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "7f1053b17c810be5",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:20:26.927643Z",
|
||||
"start_time": "2024-06-03T17:20:25.128994Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/Users/joein/work/qdrant/fastembed/venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
|
||||
" from .autonotebook import tqdm as notebook_tqdm\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "[{'model': 'colbert-ir/colbertv2.0',\n 'dim': 128,\n 'description': 'Late interaction model',\n 'size_in_GB': 0.44,\n 'sources': {'hf': 'colbert-ir/colbertv2.0'},\n 'model_file': 'model.onnx'}]"
|
||||
},
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from fastembed import LateInteractionTextEmbedding\n",
|
||||
"\n",
|
||||
"LateInteractionTextEmbedding.list_supported_models()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "c2c15893df422631",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:23:35.764183Z",
|
||||
"start_time": "2024-06-03T17:23:21.630277Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Fetching 5 files: 0%| | 0/5 [00:00<?, ?it/s]\n",
|
||||
"config.json: 100%|██████████| 743/743 [00:00<00:00, 4.56MB/s]\n",
|
||||
"\n",
|
||||
"tokenizer_config.json: 100%|██████████| 405/405 [00:00<00:00, 3.34MB/s]\n",
|
||||
"Fetching 5 files: 20%|██ | 1/5 [00:00<00:01, 3.64it/s]\n",
|
||||
"tokenizer.json: 0%| | 0.00/466k [00:00<?, ?B/s]\u001b[A\n",
|
||||
"\n",
|
||||
"special_tokens_map.json: 100%|██████████| 112/112 [00:00<00:00, 727kB/s]\n",
|
||||
"\n",
|
||||
"tokenizer.json: 100%|██████████| 466k/466k [00:00<00:00, 1.48MB/s]\u001b[A\n",
|
||||
"\n",
|
||||
"model.onnx: 0%| | 0.00/436M [00:00<?, ?B/s]\u001b[A\n",
|
||||
"model.onnx: 2%|▏ | 10.5M/436M [00:00<00:34, 12.2MB/s]\u001b[A\n",
|
||||
"model.onnx: 5%|▍ | 21.0M/436M [00:01<00:20, 20.3MB/s]\u001b[A\n",
|
||||
"model.onnx: 7%|▋ | 31.5M/436M [00:01<00:15, 25.7MB/s]\u001b[A\n",
|
||||
"model.onnx: 10%|▉ | 41.9M/436M [00:01<00:13, 29.4MB/s]\u001b[A\n",
|
||||
"model.onnx: 12%|█▏ | 52.4M/436M [00:01<00:12, 31.9MB/s]\u001b[A\n",
|
||||
"model.onnx: 14%|█▍ | 62.9M/436M [00:02<00:11, 33.7MB/s]\u001b[A\n",
|
||||
"model.onnx: 17%|█▋ | 73.4M/436M [00:02<00:10, 34.9MB/s]\u001b[A\n",
|
||||
"model.onnx: 19%|█▉ | 83.9M/436M [00:02<00:09, 35.5MB/s]\u001b[A\n",
|
||||
"model.onnx: 22%|██▏ | 94.4M/436M [00:03<00:09, 36.1MB/s]\u001b[A\n",
|
||||
"model.onnx: 24%|██▍ | 105M/436M [00:03<00:09, 36.6MB/s] \u001b[A\n",
|
||||
"model.onnx: 26%|██▋ | 115M/436M [00:03<00:08, 36.9MB/s]\u001b[A\n",
|
||||
"model.onnx: 29%|██▉ | 126M/436M [00:03<00:08, 37.1MB/s]\u001b[A\n",
|
||||
"model.onnx: 31%|███▏ | 136M/436M [00:04<00:08, 37.3MB/s]\u001b[A\n",
|
||||
"model.onnx: 34%|███▎ | 147M/436M [00:04<00:07, 37.4MB/s]\u001b[A\n",
|
||||
"model.onnx: 36%|███▌ | 157M/436M [00:04<00:07, 37.4MB/s]\u001b[A\n",
|
||||
"model.onnx: 38%|███▊ | 168M/436M [00:05<00:07, 37.5MB/s]\u001b[A\n",
|
||||
"model.onnx: 41%|████ | 178M/436M [00:05<00:06, 37.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 43%|████▎ | 189M/436M [00:05<00:06, 37.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 46%|████▌ | 199M/436M [00:05<00:06, 37.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 48%|████▊ | 210M/436M [00:06<00:06, 37.5MB/s]\u001b[A\n",
|
||||
"model.onnx: 50%|█████ | 220M/436M [00:06<00:05, 37.5MB/s]\u001b[A\n",
|
||||
"model.onnx: 53%|█████▎ | 231M/436M [00:06<00:05, 37.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 55%|█████▌ | 241M/436M [00:06<00:05, 37.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 58%|█████▊ | 252M/436M [00:07<00:04, 37.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 60%|██████ | 262M/436M [00:07<00:04, 37.7MB/s]\u001b[A\n",
|
||||
"model.onnx: 63%|██████▎ | 273M/436M [00:07<00:04, 37.7MB/s]\u001b[A\n",
|
||||
"model.onnx: 65%|██████▍ | 283M/436M [00:08<00:04, 36.0MB/s]\u001b[A\n",
|
||||
"model.onnx: 67%|██████▋ | 294M/436M [00:08<00:03, 36.4MB/s]\u001b[A\n",
|
||||
"model.onnx: 70%|██████▉ | 304M/436M [00:08<00:03, 36.8MB/s]\u001b[A\n",
|
||||
"model.onnx: 72%|███████▏ | 315M/436M [00:08<00:03, 37.0MB/s]\u001b[A\n",
|
||||
"model.onnx: 75%|███████▍ | 325M/436M [00:09<00:02, 37.3MB/s]\u001b[A\n",
|
||||
"model.onnx: 77%|███████▋ | 336M/436M [00:09<00:03, 30.8MB/s]\u001b[A\n",
|
||||
"model.onnx: 79%|███████▉ | 346M/436M [00:10<00:02, 32.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 82%|████████▏ | 357M/436M [00:10<00:02, 33.9MB/s]\u001b[A\n",
|
||||
"model.onnx: 84%|████████▍ | 367M/436M [00:10<00:01, 34.8MB/s]\u001b[A\n",
|
||||
"model.onnx: 87%|████████▋ | 377M/436M [00:10<00:01, 35.7MB/s]\u001b[A\n",
|
||||
"model.onnx: 89%|████████▉ | 388M/436M [00:11<00:01, 36.2MB/s]\u001b[A\n",
|
||||
"model.onnx: 91%|█████████▏| 398M/436M [00:11<00:01, 36.6MB/s]\u001b[A\n",
|
||||
"model.onnx: 94%|█████████▍| 409M/436M [00:11<00:00, 36.9MB/s]\u001b[A\n",
|
||||
"model.onnx: 96%|█████████▌| 419M/436M [00:11<00:00, 37.1MB/s]\u001b[A\n",
|
||||
"model.onnx: 99%|█████████▊| 430M/436M [00:12<00:00, 37.3MB/s]\u001b[A\n",
|
||||
"model.onnx: 100%|██████████| 436M/436M [00:12<00:00, 35.1MB/s]\u001b[A\n",
|
||||
"Fetching 5 files: 100%|██████████| 5/5 [00:13<00:00, 2.68s/it]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embedding_model = LateInteractionTextEmbedding(\"colbert-ir/colbertv2.0\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"id": "e560b5fa7d63bea3",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:39:33.400876Z",
|
||||
"start_time": "2024-06-03T17:39:33.397431Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"documents = [\n",
|
||||
" \"ColBERT is a late interaction text embedding model, however, there are also other models such as TwinBERT.\",\n",
|
||||
" \"On the contrary to the late interaction models, the early interaction models contains interaction steps at embedding generation process\",\n",
|
||||
"]\n",
|
||||
"queries = [\n",
|
||||
" \"Are there any other late interaction text embedding models except ColBERT?\",\n",
|
||||
" \"What is the difference between late interaction and early interaction text embedding models?\",\n",
|
||||
"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "347ad924a3449743",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"*NOTE*: ColBERT computes query and documents embeddings differently, make sure to use the corresponding methods."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 17,
|
||||
"id": "496fbf51e4eaaae",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:39:34.379885Z",
|
||||
"start_time": "2024-06-03T17:39:34.316257Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"document_embeddings = list(\n",
|
||||
" embedding_model.embed(documents)\n",
|
||||
") # embed and qury_embed return generators,\n",
|
||||
"# which we need to evaluate by writing them to a list\n",
|
||||
"query_embeddings = list(embedding_model.query_embed(queries))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"id": "50595bb0498f0c7c",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:39:34.793528Z",
|
||||
"start_time": "2024-06-03T17:39:34.788545Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "((26, 128), (32, 128))"
|
||||
},
|
||||
"execution_count": 18,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"document_embeddings[0].shape, query_embeddings[0].shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "13e43f2c24a7d5fc",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"Don't worry about query embeddings having the bigger shape in this case. \n",
|
||||
"ColBERT authors recommend to pad queries with [MASK] tokens to 32 tokens.\n",
|
||||
"They also recommends to truncate queries to 32 tokens, however we don't do that in FastEmbed, so you can put some straight into the queries."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "bb1a4011effd3699",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"## MaxSim operator"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e9ea4cf82521f2de",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"Qdrant will support ColBERT as of the next version (v1.10), however, at the moment, you can compute embedding similarities manually. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 19,
|
||||
"id": "f84392f63d2c6076",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:39:36.431622Z",
|
||||
"start_time": "2024-06-03T17:39:36.427363Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def compute_relevance_scores(query_embedding: np.array, document_embeddings: np.array, k: int):\n",
|
||||
" \"\"\"\n",
|
||||
" Compute relevance scores for top-k documents given a query.\n",
|
||||
"\n",
|
||||
" :param query_embedding: Numpy array representing the query embedding, shape: [num_query_terms, embedding_dim]\n",
|
||||
" :param document_embeddings: Numpy array representing embeddings for documents, shape: [num_documents, max_doc_length, embedding_dim]\n",
|
||||
" :param k: Number of top documents to return\n",
|
||||
" :return: Indices of the top-k documents based on their relevance scores\n",
|
||||
" \"\"\"\n",
|
||||
" # Compute batch dot-product of query_embedding and document_embeddings\n",
|
||||
" # Resulting shape: [num_documents, num_query_terms, max_doc_length]\n",
|
||||
" scores = np.matmul(query_embedding, document_embeddings.transpose(0, 2, 1))\n",
|
||||
"\n",
|
||||
" # Apply max-pooling across document terms (axis=2) to find the max similarity per query term\n",
|
||||
" # Shape after max-pool: [num_documents, num_query_terms]\n",
|
||||
" max_scores_per_query_term = np.max(scores, axis=2)\n",
|
||||
"\n",
|
||||
" # Sum the scores across query terms to get the total score for each document\n",
|
||||
" # Shape after sum: [num_documents]\n",
|
||||
" total_scores = np.sum(max_scores_per_query_term, axis=1)\n",
|
||||
"\n",
|
||||
" # Sort the documents based on their total scores and get the indices of the top-k documents\n",
|
||||
" sorted_indices = np.argsort(total_scores)[::-1][:k]\n",
|
||||
"\n",
|
||||
" return sorted_indices"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"id": "c61d07bed7b60e35",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:39:37.053383Z",
|
||||
"start_time": "2024-06-03T17:39:37.050926Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Sorted document indices: [0 1]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"sorted_indices = compute_relevance_scores(\n",
|
||||
" np.array(query_embeddings[0]), np.array(document_embeddings), k=3\n",
|
||||
")\n",
|
||||
"print(\"Sorted document indices:\", sorted_indices)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 22,
|
||||
"id": "b24df2569970d9e8",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-03T17:40:52.276846Z",
|
||||
"start_time": "2024-06-03T17:40:52.273789Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Query: Are there any other late interaction text embedding models except ColBERT?\n",
|
||||
"Document: ColBERT is a late interaction text embedding model, however, there are also other models such as TwinBERT.\n",
|
||||
"Document: On the contrary to the late interaction models, the early interaction models contains interaction steps at embedding generation process\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(f\"Query: {queries[0]}\")\n",
|
||||
"for index in sorted_indices:\n",
|
||||
" print(f\"Document: {documents[index]}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6de537c37aff3927",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"## Use-case recommendation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "37e3525d3259cd2b",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"Despite ColBERT allows to compute embeddings independently and spare some workload offline, it still computes more resources than no interaction models. Due to this, it might be more reasonable to use ColBERT not as a first-stage retriever, but as a re-ranker.\n",
|
||||
"\n",
|
||||
"The first-stage retriever would then be a no-interaction model, which e.g. retrieves first 100 or 500 examples, and leave the final ranking to the ColBERT model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "cfa922793454b4ad",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython2",
|
||||
"version": "2.7.6"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -1,436 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ntGNDuSCeAR2"
|
||||
},
|
||||
"source": [
|
||||
"# FastEmbed on GPU\n",
|
||||
"\n",
|
||||
"As of version 0.2.7 FastEmbed supports GPU acceleration.\n",
|
||||
"\n",
|
||||
"This notebook covers the installation process and usage of fastembed on GPU.\n",
|
||||
"\n",
|
||||
"## Installation\n",
|
||||
"\n",
|
||||
"Fastembed depends on `onnxruntime` and inherits its scheme of GPU support.\n",
|
||||
"\n",
|
||||
"In order to use GPU with onnx models, you would need to have `onnxruntime-gpu` package, which substitutes all the `onnxruntime` functionality.\n",
|
||||
"Fastembed mimics this behavior and requires `fastembed-gpu` package to be installed."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "GK2XADwUeEK7"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install fastembed-gpu"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "3aiGPqjCeGzo"
|
||||
},
|
||||
"source": [
|
||||
"**NOTE**: `onnxruntime-gpu` and `onnxruntime` can't be installed in the same environment. If you have `onnxruntime` installed, you would need to uninstall it before installing `onnxruntime-gpu`. Same is true for `fastembed` and `fastembed-gpu`.\n",
|
||||
"\n",
|
||||
"### CUDA 12.x support\n",
|
||||
"\n",
|
||||
"By default `onnxruntime-gpu` is shipped with CUDA 11.8 support.\n",
|
||||
"CUDA 12.x support requires installation of `onnxruntime-gpu` with providing of a direct url:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "OoSfWFFZeJ5t",
|
||||
"outputId": "417b9332-6a7b-4000-c74b-4ed2b5b76590"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install onnxruntime-gpu -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ -qq\n",
|
||||
"!pip install fastembed-gpu -qqq"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "3xx3r-9jgAMi"
|
||||
},
|
||||
"source": [
|
||||
"You can check your CUDA version using such commands as `nvidia-smi` or `nvcc --version`\n",
|
||||
"\n",
|
||||
"Google Colab notebooks have CUDA 12.x."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Igv5RXhSeO68"
|
||||
},
|
||||
"source": [
|
||||
"### CUDA drivers\n",
|
||||
"\n",
|
||||
"FastEmbed does not include CUDA drivers and CuDNN libraries.\n",
|
||||
"You would need to take care of the environment setup on your own.\n",
|
||||
"Dependencies required for the chosen onnxruntime version can be found [here](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Usage"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 334,
|
||||
"referenced_widgets": [
|
||||
"aacf08a7aa444b64a2efad1967d28a53",
|
||||
"5606aa785de74d65a9928b31c0be8a53",
|
||||
"d4ec9d3b74ec4412894da2161ed2bddf",
|
||||
"8edd544c3e074ec1813e5b9d1aef43d9",
|
||||
"9898890f8a75468ea20e3ce319d0b6e2",
|
||||
"da3b18abb16241a0a7191ee9afcb0510",
|
||||
"258a619168824253a6a329efdc51ebe6",
|
||||
"53c7cdc967d24faba0b5c659c94c50b8",
|
||||
"e9348d8be28d408e8e760c71b21ab294",
|
||||
"0ba06e0816714f2fbdec8260f160abc0",
|
||||
"0b96563334964d449dd34f35b6b3e715",
|
||||
"11c2eec490e8479b944eec7f30cb1ca2",
|
||||
"91463da0d1c5466795e06ab586002259",
|
||||
"30f4f7833406474f89ef0700b00a33aa",
|
||||
"4302c304ec6a4b5985797e300bd7e353",
|
||||
"2605640c7b824ed7aa137d404e14b774",
|
||||
"b02efe3a33d04f06aa8938719ab35671",
|
||||
"50408e5d052343b1a1b44a0fae0f801d",
|
||||
"1be01c95d9e84f8ea88367c987a72fdc",
|
||||
"a109c13bc93a449186424542dc330be8",
|
||||
"4adce304ce1947b5a01dde10bbb3bb8c",
|
||||
"a761366a37e44837a25e0f25b18efed2",
|
||||
"94512b9055e546389471197b76ad5449",
|
||||
"072dca00bd7b4918a178f90ccabf698a",
|
||||
"48a856c59ef74cc3834521b1bf616541",
|
||||
"c020c503aeaa464cad643ade5ee3ae24",
|
||||
"a4e7e40c0bbd4f878c20a9f65fe3a048",
|
||||
"e8c0a1c339fd47668d944a9defad79d4",
|
||||
"cd782d35c6bd40c0a60d57b1828a7251",
|
||||
"04f638ab08da4d20928644c4ba03f8ef",
|
||||
"17f20477fc79475f97adf1c1f64a4192",
|
||||
"96f7b5a2e224462e9fcffd03f906a593",
|
||||
"755cd32d9fc9407c80a160f45c802d1e",
|
||||
"a886258e7cd14c048b58391d7b772901",
|
||||
"bc3e48f826a74840867a6209e622b75e",
|
||||
"125b2ac0f78043bba7eca53474ca44c4",
|
||||
"82f186d1ffb4435d94a6c7e9025242ef",
|
||||
"77000333e5ca4094be291ad82d4a627a",
|
||||
"7fe64fb53055431488d002c76c8e331e",
|
||||
"7de59ae9919f4a5bb2b6e601a3c02412",
|
||||
"97a69423a6644eab87fc636e182f23a4",
|
||||
"4df936d1065b41f4bf02ed394fdf7b7e",
|
||||
"3918bd1affa3454e8e9044a418a056ea",
|
||||
"163b27ae0bce41e5b48efcb4b3fd780d",
|
||||
"94631fd6e0744085bc79c3121de4a9f7",
|
||||
"31cd98d66bc54418b35e70fbbc0fa3c0",
|
||||
"6d21627a638b4ddca6fe7bfb80a621b5",
|
||||
"b37bed9dc4fe45c08b8397288fe5b1a9",
|
||||
"164fef95d1414177a40d563f5682f6a3",
|
||||
"1a9a0ea53448413a8e4b360b7bb69e26",
|
||||
"dd1a4483b4b045c6929e3d2cf1338f63",
|
||||
"496ddd8e05f949cd8cbba8e677f476ac",
|
||||
"2813be951d7f48b2aad1dd4a444ce3eb",
|
||||
"8e9a2c2dd21942edbdfecb3b7dffc70b",
|
||||
"08a10fe247f1425db044cfc13f2fb384",
|
||||
"b8786aded92d421592bc7623c5c7899e",
|
||||
"c91a20a9433d4016ba2db69fa50e0b4d",
|
||||
"e997820738594c6dadb061908d7afdc1",
|
||||
"a5fc751f81ae498f9aa55ece0e6853b2",
|
||||
"2aee4fc8cda64c5eb8722be81e48e0ca",
|
||||
"3a53e8624dff48b3959875ef58ee99ce",
|
||||
"50a70044f77542108fe188598e70797e",
|
||||
"13cf998b35ae4507a63e797f6fa3eada",
|
||||
"6209eb6a68cf4a378767ef34d0d9216d",
|
||||
"7395db766b944af9b41d6b56c9ada0b1",
|
||||
"42122c317ec648688f0164a1adb5df28"
|
||||
]
|
||||
},
|
||||
"id": "Ttf4YggPeQQK",
|
||||
"outputId": "aa75129d-9e2d-4c88-cf03-251dd43a11b1"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:88: UserWarning: \n",
|
||||
"The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
|
||||
"To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
|
||||
"You will be able to reuse this secret in all of your notebooks.\n",
|
||||
"Please note that authentication is recommended but still optional to access public models or datasets.\n",
|
||||
" warnings.warn(\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "aacf08a7aa444b64a2efad1967d28a53",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Fetching 5 files: 0%| | 0/5 [00:00<?, ?it/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "11c2eec490e8479b944eec7f30cb1ca2",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"tokenizer_config.json: 0%| | 0.00/1.24k [00:00<?, ?B/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "94512b9055e546389471197b76ad5449",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"config.json: 0%| | 0.00/706 [00:00<?, ?B/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "a886258e7cd14c048b58391d7b772901",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"special_tokens_map.json: 0%| | 0.00/695 [00:00<?, ?B/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "94631fd6e0744085bc79c3121de4a9f7",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"tokenizer.json: 0%| | 0.00/711k [00:00<?, ?B/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "b8786aded92d421592bc7623c5c7899e",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"model_optimized.onnx: 0%| | 0.00/66.5M [00:00<?, ?B/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"['CUDAExecutionProvider', 'CPUExecutionProvider']"
|
||||
]
|
||||
},
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from typing import List\n",
|
||||
"\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"from fastembed import TextEmbedding\n",
|
||||
"\n",
|
||||
"embedding_model_gpu = TextEmbedding(\n",
|
||||
" model_name=\"BAAI/bge-small-en-v1.5\", providers=[\"CUDAExecutionProvider\"]\n",
|
||||
")\n",
|
||||
"embedding_model_gpu.model.model.get_providers()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"id": "iPtoHf7GeV-i"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"documents: List[str] = list(np.repeat(\"Demonstrating GPU acceleration in fastembed\", 500))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "islhyLf4ed-H",
|
||||
"outputId": "8c8ed09b-9eac-438f-97bc-578751975148"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"43.4 ms ± 2.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%%timeit\n",
|
||||
"list(embedding_model_gpu.embed(documents))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 67,
|
||||
"referenced_widgets": [
|
||||
"9c306ce5188c45feb8dfb9089592591c",
|
||||
"296ff54c6e61441f978084df59626598",
|
||||
"d6d42b4f245a49b7ba7769e23a3202fc",
|
||||
"39ce7754480147759c16a3089d8105af",
|
||||
"8253960a069d4106863a75faae54b90d",
|
||||
"7ccf959452af4c0b873c7567747f0816",
|
||||
"ac9d0b5a5b1f401e90a1cc9ffe6d4b4c",
|
||||
"0aada067dec3472f9aba1772d6b775a5",
|
||||
"07597b1287e04653b80c47a771549376",
|
||||
"054be1dd9f084cae911745b692ccd929",
|
||||
"ab19e8e831694e308a4b79f05aff728e"
|
||||
]
|
||||
},
|
||||
"id": "bOKVUvWJegYJ",
|
||||
"outputId": "dde74917-08b0-4ce2-9a2b-cc31e02cafb2"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "9c306ce5188c45feb8dfb9089592591c",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Fetching 5 files: 0%| | 0/5 [00:00<?, ?it/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"['CPUExecutionProvider']"
|
||||
]
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embedding_model_cpu = TextEmbedding(model_name=\"BAAI/bge-small-en-v1.5\")\n",
|
||||
"embedding_model_cpu.model.model.get_providers()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "0NJj9RvSfASP",
|
||||
"outputId": "526f5280-99bd-454e-8af8-6a860ad96e54"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"4.33 s ± 591 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%%timeit\n",
|
||||
"list(embedding_model_cpu.embed(documents))"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": []
|
||||
},
|
||||
"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.12"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 1
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,128 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "aa0a86859809102",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"# Image Embedding\n",
|
||||
"As of version 0.3.0 fastembed supports computation of image embeddings.\n",
|
||||
"\n",
|
||||
"The process is as easy and straightforward as with text embeddings. Let's see how it works."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "cea8fd5c019571fe",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-02T11:35:40.126023Z",
|
||||
"start_time": "2024-06-02T11:35:39.864701Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Fetching 3 files: 100%|██████████| 3/3 [00:00<00:00, 47482.69it/s]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "[array([0. , 0. , 0. , ..., 0. , 0.01139933,\n 0. ], dtype=float32),\n array([0.02169187, 0. , 0. , ..., 0. , 0.00848291,\n 0. ], dtype=float32)]"
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from fastembed import ImageEmbedding\n",
|
||||
"\n",
|
||||
"model = ImageEmbedding(\"Qdrant/resnet50-onnx\")\n",
|
||||
"\n",
|
||||
"embeddings_generator = model.embed(\n",
|
||||
" [\"../../tests/misc/image.jpeg\", \"../../tests/misc/small_image.jpeg\"]\n",
|
||||
")\n",
|
||||
"embeddings_list = list(embeddings_generator)\n",
|
||||
"embeddings_list"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3f838f18523ad1e0",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"## Preprocessing\n",
|
||||
"\n",
|
||||
"Preprocessing is encapsulated in the ImageEmbedding class, applied operations are identical to the ones provided by [Hugging Face Transformers](https://huggingface.co/docs/transformers/en/index).\n",
|
||||
"You don't need to think about batching, opening/closing files, resizing images, etc., Fastembed will take care of it."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "894b33ff9b385d72",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"## Supported models\n",
|
||||
"\n",
|
||||
"List of supported image embedding models can either be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-image-embedding-models) or by calling the `ImageEmbedding.list_supported_models()` method."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "6d6a4cbbd2200d14",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-02T11:40:19.313226Z",
|
||||
"start_time": "2024-06-02T11:40:19.309845Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "[{'model': 'Qdrant/clip-ViT-B-32-vision',\n 'dim': 512,\n 'description': 'CLIP vision encoder based on ViT-B/32',\n 'size_in_GB': 0.34,\n 'sources': {'hf': 'Qdrant/clip-ViT-B-32-vision'},\n 'model_file': 'model.onnx'},\n {'model': 'Qdrant/resnet50-onnx',\n 'dim': 2048,\n 'description': 'ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`__.',\n 'size_in_GB': 0.1,\n 'sources': {'hf': 'Qdrant/resnet50-onnx'},\n 'model_file': 'model.onnx'}]"
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"ImageEmbedding.list_supported_models()"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython2",
|
||||
"version": "2.7.6"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
+12
-4
@@ -39,7 +39,15 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32m2024-02-07 22:20:57.013\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mfastembed.embedding\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m7\u001b[0m - \u001b[33m\u001b[1mDefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated. Use TextEmbedding instead.\u001b[0m\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from typing import List\n",
|
||||
"import numpy as np\n",
|
||||
@@ -84,7 +92,7 @@
|
||||
" \"His life has been depicted in various films, TV shows, and books\",\n",
|
||||
"]\n",
|
||||
"# Initialize the DefaultEmbedding class with the desired parameters\n",
|
||||
"embedding_model = TextEmbedding(model_name=\"BAAI/bge-small-en\")\n",
|
||||
"embedding_model = TextEmbedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
|
||||
"\n",
|
||||
"# We'll use the passage_embed method to get the embeddings for the documents\n",
|
||||
"embeddings: List[np.ndarray] = list(\n",
|
||||
@@ -132,7 +140,7 @@
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"(array([-0.06002192, 0.04322132, -0.00545516, -0.04419701, -0.00542277],\n",
|
||||
"(array([-0.04393955, 0.04452892, -0.00760788, -0.03399807, 0.01951348],\n",
|
||||
" dtype=float32),\n",
|
||||
" array([-0.06002192, 0.04322132, -0.00545516, -0.04419701, -0.00542277],\n",
|
||||
" dtype=float32))"
|
||||
@@ -175,7 +183,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.13"
|
||||
"version": "3.11.5"
|
||||
},
|
||||
"orig_nbformat": 4
|
||||
},
|
||||
@@ -1,384 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Introduction to SPLADE with FastEmbed\n",
|
||||
"\n",
|
||||
"In this notebook, we will explore how to generate Sparse Vectors -- in particular a variant of the [SPLADE](https://arxiv.org/abs/2107.05720).\n",
|
||||
"\n",
|
||||
"> 💡 The original [naver/SPLADE](https://github.com/naver/splade) models were licensed CC BY-NC-SA 4.0 -- Not for Commercial Use. This [SPLADE++](https://huggingface.co/prithivida/Splade_PP_en_v1) model is Apache License and hence, licensed for commercial use. \n",
|
||||
"\n",
|
||||
"## Outline:\n",
|
||||
"1. [What is SPLADE?](#What-is-SPLADE?)\n",
|
||||
"2. [Setting up the environment](#Setting-up-the-environment)\n",
|
||||
"3. [Generating SPLADE vectors with FastEmbed](#Generating-SPLADE-vectors-with-FastEmbed)\n",
|
||||
"4. [Understanding SPLADE vectors](#Understanding-SPLADE-vectors)\n",
|
||||
"5. [Observations and Design Choices](#Observations-and-Model-Design-Choices)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"## What is SPLADE?\n",
|
||||
"\n",
|
||||
"SPLADE was a novel method for _learning_ sparse vectors for text representation. This model beats BM25 -- the underlying approach for the Elastic/Lucene family of implementations. Thus making it highly effective for tasks such as information retrieval, document classification, and more. \n",
|
||||
"\n",
|
||||
"The key advantage of SPLADE is its ability to generate sparse vectors, which are more efficient and interpretable than dense vectors. This makes SPLADE a powerful tool for handling large-scale text data.\n",
|
||||
"\n",
|
||||
"## Setting up the environment\n",
|
||||
"\n",
|
||||
"This notebook uses few dependencies, which are installed below: "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# !pip install -q fastembed"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Let's get started! 🚀"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:20.516644Z",
|
||||
"start_time": "2024-03-30T00:49:20.188543Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from fastembed import SparseTextEmbedding, SparseEmbedding\n",
|
||||
"from typing import List"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"> You can find the list of all supported Sparse Embedding models by calling this API: `SparseTextEmbedding.list_supported_models()`"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:22.366294Z",
|
||||
"start_time": "2024-03-30T00:49:22.362384Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'model': 'prithvida/Splade_PP_en_v1',\n",
|
||||
" 'vocab_size': 30522,\n",
|
||||
" 'description': 'Misspelled version of the model. Retained for backward compatibility. Independent Implementation of SPLADE++ Model for English',\n",
|
||||
" 'size_in_GB': 0.532,\n",
|
||||
" 'sources': {'hf': 'Qdrant/SPLADE_PP_en_v1'}},\n",
|
||||
" {'model': 'prithivida/Splade_PP_en_v1',\n",
|
||||
" 'vocab_size': 30522,\n",
|
||||
" 'description': 'Independent Implementation of SPLADE++ Model for English',\n",
|
||||
" 'size_in_GB': 0.532,\n",
|
||||
" 'sources': {'hf': 'Qdrant/SPLADE_PP_en_v1'}}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"SparseTextEmbedding.list_supported_models()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:27.193530Z",
|
||||
"start_time": "2024-03-30T00:49:26.139248Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "2aa47b26ab01475e8d3577433037f685",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Fetching 9 files: 0%| | 0/9 [00:00<?, ?it/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"model_name = \"prithvida/Splade_PP_en_v1\"\n",
|
||||
"# This triggers the model download\n",
|
||||
"model = SparseTextEmbedding(model_name=model_name)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:28.624109Z",
|
||||
"start_time": "2024-03-30T00:49:28.399960Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"documents: List[str] = [\n",
|
||||
" \"Chandrayaan-3 is India's third lunar mission\",\n",
|
||||
" \"It aimed to land a rover on the Moon's surface - joining the US, China and Russia\",\n",
|
||||
" \"The mission is a follow-up to Chandrayaan-2, which had partial success\",\n",
|
||||
" \"Chandrayaan-3 will be launched by the Indian Space Research Organisation (ISRO)\",\n",
|
||||
" \"The estimated cost of the mission is around $35 million\",\n",
|
||||
" \"It will carry instruments to study the lunar surface and atmosphere\",\n",
|
||||
" \"Chandrayaan-3 landed on the Moon's surface on 23rd August 2023\",\n",
|
||||
" \"It consists of a lander named Vikram and a rover named Pragyan similar to Chandrayaan-2. Its propulsion module would act like an orbiter.\",\n",
|
||||
" \"The propulsion module carries the lander and rover configuration until the spacecraft is in a 100-kilometre (62 mi) lunar orbit\",\n",
|
||||
" \"The mission used GSLV Mk III rocket for its launch\",\n",
|
||||
" \"Chandrayaan-3 was launched from the Satish Dhawan Space Centre in Sriharikota\",\n",
|
||||
" \"Chandrayaan-3 was launched earlier in the year 2023\",\n",
|
||||
"]\n",
|
||||
"sparse_embeddings_list: List[SparseEmbedding] = list(\n",
|
||||
" model.embed(documents, batch_size=6)\n",
|
||||
") # batch_size is optional, notice the generator"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:29.646340Z",
|
||||
"start_time": "2024-03-30T00:49:29.643411Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"SparseEmbedding(values=array([0.05297208, 0.01963477, 0.36459631, 1.38508618, 0.71776593,\n",
|
||||
" 0.12667948, 0.46230844, 0.446771 , 0.26897505, 1.01519883,\n",
|
||||
" 1.5655334 , 0.29412213, 1.53102326, 0.59785569, 1.1001817 ,\n",
|
||||
" 0.02079751, 0.09955651, 0.44249091, 0.09747757, 1.53519952,\n",
|
||||
" 1.36765671, 0.15740395, 0.49882549, 0.38629025, 0.76612782,\n",
|
||||
" 1.25805044, 0.39058095, 0.27236196, 0.45152301, 0.48262018,\n",
|
||||
" 0.26085234, 1.35912788, 0.70710695, 1.71639752]), indices=array([ 1010, 1011, 1016, 1017, 2001, 2018, 2034, 2093, 2117,\n",
|
||||
" 2319, 2353, 2509, 2634, 2686, 2796, 2817, 2922, 2959,\n",
|
||||
" 3003, 3148, 3260, 3390, 3462, 3523, 3822, 4231, 4316,\n",
|
||||
" 4774, 5590, 5871, 6416, 11926, 12076, 16469]))"
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"index = 0\n",
|
||||
"sparse_embeddings_list[index]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The previous output is a SparseEmbedding object for the first document in our list.\n",
|
||||
"\n",
|
||||
"It contains two arrays: values and indices. \n",
|
||||
"- The 'values' array represents the weights of the features (tokens) in the document.\n",
|
||||
"- The 'indices' array represents the indices of these features in the model's vocabulary.\n",
|
||||
"\n",
|
||||
"Each pair of corresponding values and indices represents a token and its weight in the document."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:31.549533Z",
|
||||
"start_time": "2024-03-30T00:49:31.546398Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Token at index 1010 has weight 0.05297207832336426\n",
|
||||
"Token at index 1011 has weight 0.01963476650416851\n",
|
||||
"Token at index 1016 has weight 0.36459630727767944\n",
|
||||
"Token at index 1017 has weight 1.385086178779602\n",
|
||||
"Token at index 2001 has weight 0.7177659273147583\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Let's print the first 5 features and their weights for better understanding.\n",
|
||||
"for i in range(5):\n",
|
||||
" print(f\"Token at index {sparse_embeddings_list[0].indices[i]} has weight {sparse_embeddings_list[0].values[i]}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Understanding SPLADE vectors\n",
|
||||
"\n",
|
||||
"This is still a little abstract, so let's use the tokenizer vocab to make sense of these indices."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:36.203640Z",
|
||||
"start_time": "2024-03-30T00:49:34.889654Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import json\n",
|
||||
"from transformers import AutoTokenizer\n",
|
||||
"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(SparseTextEmbedding.list_supported_models()[0][\"sources\"][\"hf\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-03-30T00:49:36.210049Z",
|
||||
"start_time": "2024-03-30T00:49:36.206825Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"chandra\": 1.7163975238800049,\n",
|
||||
" \"third\": 1.5655333995819092,\n",
|
||||
" \"##ya\": 1.535199522972107,\n",
|
||||
" \"india\": 1.5310232639312744,\n",
|
||||
" \"3\": 1.385086178779602,\n",
|
||||
" \"mission\": 1.3676567077636719,\n",
|
||||
" \"lunar\": 1.3591278791427612,\n",
|
||||
" \"moon\": 1.2580504417419434,\n",
|
||||
" \"indian\": 1.1001816987991333,\n",
|
||||
" \"##an\": 1.015198826789856,\n",
|
||||
" \"3rd\": 0.7661278247833252,\n",
|
||||
" \"was\": 0.7177659273147583,\n",
|
||||
" \"spacecraft\": 0.7071069478988647,\n",
|
||||
" \"space\": 0.5978556871414185,\n",
|
||||
" \"flight\": 0.4988254904747009,\n",
|
||||
" \"satellite\": 0.4826201796531677,\n",
|
||||
" \"first\": 0.46230843663215637,\n",
|
||||
" \"expedition\": 0.4515230059623718,\n",
|
||||
" \"three\": 0.4467709958553314,\n",
|
||||
" \"fourth\": 0.44249090552330017,\n",
|
||||
" \"vehicle\": 0.390580952167511,\n",
|
||||
" \"iii\": 0.3862902522087097,\n",
|
||||
" \"2\": 0.36459630727767944,\n",
|
||||
" \"##3\": 0.2941221296787262,\n",
|
||||
" \"planet\": 0.27236196398735046,\n",
|
||||
" \"second\": 0.26897504925727844,\n",
|
||||
" \"missions\": 0.2608523368835449,\n",
|
||||
" \"launched\": 0.15740394592285156,\n",
|
||||
" \"had\": 0.12667948007583618,\n",
|
||||
" \"largest\": 0.09955651313066483,\n",
|
||||
" \"leader\": 0.09747757017612457,\n",
|
||||
" \",\": 0.05297207832336426,\n",
|
||||
" \"study\": 0.02079751156270504,\n",
|
||||
" \"-\": 0.01963476650416851\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"def get_tokens_and_weights(sparse_embedding, tokenizer):\n",
|
||||
" token_weight_dict = {}\n",
|
||||
" for i in range(len(sparse_embedding.indices)):\n",
|
||||
" token = tokenizer.decode([sparse_embedding.indices[i]])\n",
|
||||
" weight = sparse_embedding.values[i]\n",
|
||||
" token_weight_dict[token] = weight\n",
|
||||
"\n",
|
||||
" # Sort the dictionary by weights\n",
|
||||
" token_weight_dict = dict(sorted(token_weight_dict.items(), key=lambda item: item[1], reverse=True))\n",
|
||||
" return token_weight_dict\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# Test the function with the first SparseEmbedding\n",
|
||||
"print(json.dumps(get_tokens_and_weights(sparse_embeddings_list[index], tokenizer), indent=4))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Observations and Model Design Choices\n",
|
||||
"\n",
|
||||
"1. The relative order of importance is quite useful. The most important tokens in the sentence have the highest weights.\n",
|
||||
"1. **Term Expansion**: The model can expand the terms in the document. This means that the model can generate weights for tokens that are not present in the document but are related to the tokens in the document. This is a powerful feature that allows the model to capture the context of the document. Here, you'll see that the model has added the tokens '3' from 'third' and 'moon' from 'lunar' to the sparse vector.\n",
|
||||
"\n",
|
||||
"### Design Choices\n",
|
||||
"\n",
|
||||
"1. The weights are not normalized. This means that the sum of the weights is not 1 or 100. This is a common practice in sparse embeddings, as it allows the model to capture the importance of each token in the document.\n",
|
||||
"1. Tokens are included in the sparse vector only if they are present in the model's vocabulary. This means that the model will not generate a weight for tokens that it has not seen during training.\n",
|
||||
"1. Tokens do not map to words directly -- allowing you to gracefully handle typo errors and out-of-vocabulary tokens."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "fst",
|
||||
"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"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -3,7 +3,22 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": []
|
||||
"source": [
|
||||
"# Binary Quantization of OpenAI Embedding\n",
|
||||
"---\n",
|
||||
"\n",
|
||||
"In the world of large-scale data retrieval and processing, efficiency is crucial. With the exponential growth of data, the ability to retrieve information quickly and accurately can significantly affect system performance. This blog post explores a technique known as binary quantization applied to OpenAI embeddings, demonstrating how it can enhance **retrieval latency by 20x** or more.\n",
|
||||
"\n",
|
||||
"## What Are OpenAI Embeddings?\n",
|
||||
"OpenAI embeddings are numerical representations of textual information. They transform text into a vector space where semantically similar texts are mapped close together. This mathematical representation enables computers to understand and process human language more effectively.\n",
|
||||
"\n",
|
||||
"## Binary Quantization\n",
|
||||
"Binary quantization is a method which converts continuous numerical values into binary values (0 or 1). It simplifies the data structure, allowing faster computations. Here's a brief overview of the binary quantization process applied to OpenAI embeddings:\n",
|
||||
"\n",
|
||||
"1. **Load Embeddings**: OpenAI embeddings are loaded from parquet files.\n",
|
||||
"2. **Binary Transformation**: The continuous valued vectors are converted into binary form. Here, values greater than 0 are set to 1, and others remain 0.\n",
|
||||
"3. **Comparison & Retrieval**: Binary vectors are used for comparison using logical XOR operations and other efficient algorithms."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+12
-11
@@ -2,17 +2,17 @@
|
||||
|
||||
FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a Github issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
|
||||
|
||||
The default embedding supports "query" and "passage" prefixes for the input text. The default model is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
|
||||
|
||||
1. Light & Fast
|
||||
- Quantized model weights
|
||||
- ONNX Runtime for inference
|
||||
- ONNX Runtime for inference via [Optimum](github.com/huggingface/optimum)
|
||||
|
||||
2. Accuracy/Recall
|
||||
- Better than OpenAI Ada-002
|
||||
- Default is Flag Embedding, which has shown good results on the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard
|
||||
- Default is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard
|
||||
- List of [supported models](https://qdrant.github.io/fastembed/examples/Supported_Models/) - including multilingual models
|
||||
|
||||
Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
To install the FastEmbed library, pip works:
|
||||
@@ -24,16 +24,16 @@ pip install fastembed
|
||||
## 📖 Usage
|
||||
|
||||
```python
|
||||
from fastembed import TextEmbedding
|
||||
from fastembed.embedding import FlagEmbedding as Embedding
|
||||
|
||||
documents: List[str] = [
|
||||
"passage: Hello, World!",
|
||||
"query: Hello, World!",
|
||||
"query: Hello, World!", # these are two different embedding
|
||||
"passage: This is an example passage.",
|
||||
"fastembed is supported by and maintained by Qdrant."
|
||||
"fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
|
||||
]
|
||||
embedding_model = TextEmbedding()
|
||||
embeddings: List[np.ndarray] = embedding_model.embed(documents)
|
||||
embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512)
|
||||
embeddings: List[np.ndarray] = embedding_model.embed(documents) # If you use
|
||||
```
|
||||
|
||||
## Usage with Qdrant
|
||||
@@ -50,16 +50,17 @@ Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
|
||||
from qdrant_client import QdrantClient
|
||||
|
||||
# Initialize the client
|
||||
client = QdrantClient(":memory:") # Using an in-process Qdrant
|
||||
client = QdrantClient(":memory:") # or QdrantClient(path="path/to/db")
|
||||
|
||||
# Prepare your documents, metadata, and IDs
|
||||
docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
|
||||
metadata = [
|
||||
{"source": "Langchain-docs"},
|
||||
{"source": "Llama-index-docs"},
|
||||
{"source": "Linkedin-docs"},
|
||||
]
|
||||
ids = [42, 2]
|
||||
|
||||
# Use the new add method
|
||||
client.add(
|
||||
collection_name="demo_collection",
|
||||
documents=docs,
|
||||
|
||||
@@ -1,377 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Binary Quantization with Qdrant & OpenAI Embedding\n",
|
||||
"\n",
|
||||
"---\n",
|
||||
"In the world of large-scale data retrieval and processing, efficiency is crucial. With the exponential growth of data, the ability to retrieve information quickly and accurately can significantly affect system performance. This blog post explores a technique known as binary quantization applied to OpenAI embeddings, demonstrating how it can enhance **retrieval latency by 20x** or more.\n",
|
||||
"\n",
|
||||
"## What Are OpenAI Embeddings?\n",
|
||||
"OpenAI embeddings are numerical representations of textual information. They transform text into a vector space where semantically similar texts are mapped close together. This mathematical representation enables computers to understand and process human language more effectively.\n",
|
||||
"\n",
|
||||
"## Binary Quantization\n",
|
||||
"Binary quantization is a method which converts continuous numerical values into binary values (0 or 1). It simplifies the data structure, allowing faster computations. Here's a brief overview of the binary quantization process applied to OpenAI embeddings:\n",
|
||||
"\n",
|
||||
"1. **Load Embeddings**: OpenAI embeddings are loaded from parquet files.\n",
|
||||
"2. **Binary Transformation**: The continuous valued vectors are converted into binary form. Here, values greater than 0 are set to 1, and others remain 0.\n",
|
||||
"3. **Comparison & Retrieval**: Binary vectors are used for comparison using logical XOR operations and other efficient algorithms.\n",
|
||||
"\n",
|
||||
"Binary Quantization is a promising approach to improve retrieval speeds and reduce memory footprint of vector search engines. In this notebook we will show how to use Qdrant to perform binary quantization of vectors and perform fast similarity search on the resulting index.\n",
|
||||
"\n",
|
||||
"## Table of Contents\n",
|
||||
"1. Imports\n",
|
||||
"2. Download and Slice Dataset\n",
|
||||
"3. Create Qdrant Collection\n",
|
||||
"4. Indexing\n",
|
||||
"5. Search\n",
|
||||
"\n",
|
||||
"## 1. Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:33:03.013948Z",
|
||||
"start_time": "2024-04-01T16:33:01.019043Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install qdrant-client pandas dataset --quiet --upgrade"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:33:03.914729Z",
|
||||
"start_time": "2024-04-01T16:33:03.015394Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/Users/joein/work/qdrant/fastembed/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
|
||||
" from .autonotebook import tqdm as notebook_tqdm\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import random\n",
|
||||
"import time\n",
|
||||
"\n",
|
||||
"import numpy as np\n",
|
||||
"import pandas as pd\n",
|
||||
"from qdrant_client import QdrantClient, models\n",
|
||||
"\n",
|
||||
"random.seed(37)\n",
|
||||
"np.random.seed(37)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. Download and Slice Dataset\n",
|
||||
"\n",
|
||||
"We will be using the [dbpedia-entities](https://huggingface.co/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K) dataset from the [HuggingFace Datasets](https://huggingface.co/datasets) library. This contains 100K vectors of 1536 dimensions each"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:33:09.085853Z",
|
||||
"start_time": "2024-04-01T16:33:03.912688Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "100000"
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import datasets\n",
|
||||
"\n",
|
||||
"dataset = datasets.load_dataset(\n",
|
||||
" \"Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K\", split=\"train\"\n",
|
||||
")\n",
|
||||
"len(dataset)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:33:09.176212Z",
|
||||
"start_time": "2024-04-01T16:33:09.084550Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "True"
|
||||
},
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"client = QdrantClient(\n",
|
||||
" prefer_grpc=True,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"collection_name = \"binary-quantization\"\n",
|
||||
"client.recreate_collection(\n",
|
||||
" collection_name=collection_name,\n",
|
||||
" vectors_config=models.VectorParams(\n",
|
||||
" size=1536,\n",
|
||||
" distance=models.Distance.DOT,\n",
|
||||
" on_disk=True,\n",
|
||||
" ),\n",
|
||||
" quantization_config=models.BinaryQuantization(\n",
|
||||
" binary=models.BinaryQuantizationConfig(always_ram=True),\n",
|
||||
" ),\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:34:13.383986Z",
|
||||
"start_time": "2024-04-01T16:33:09.175725Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def iter_dataset(dataset):\n",
|
||||
" for point in dataset:\n",
|
||||
" yield point[\"openai\"], {\"text\": point[\"text\"]}\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"vectors, payload = zip(*iter_dataset(dataset))\n",
|
||||
"client.upload_collection(\n",
|
||||
" collection_name=collection_name,\n",
|
||||
" vectors=vectors,\n",
|
||||
" payload=payload,\n",
|
||||
" parallel=max(1, (os.cpu_count() // 2)),\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:34:13.390886Z",
|
||||
"start_time": "2024-04-01T16:34:13.385961Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "{'status': <CollectionStatus.YELLOW: 'yellow'>,\n 'optimizer_status': <OptimizersStatusOneOf.OK: 'ok'>,\n 'vectors_count': 116640,\n 'indexed_vectors_count': 43520,\n 'points_count': 116640,\n 'segments_count': 6,\n 'config': {'params': {'vectors': {'size': 1536,\n 'distance': <Distance.DOT: 'Dot'>,\n 'hnsw_config': None,\n 'quantization_config': None,\n 'on_disk': True},\n 'shard_number': 1,\n 'sharding_method': None,\n 'replication_factor': 1,\n 'write_consistency_factor': 1,\n 'read_fan_out_factor': None,\n 'on_disk_payload': True,\n 'sparse_vectors': None},\n 'hnsw_config': {'m': 16,\n 'ef_construct': 100,\n 'full_scan_threshold': 10000,\n 'max_indexing_threads': 0,\n 'on_disk': False,\n 'payload_m': None},\n 'optimizer_config': {'deleted_threshold': 0.2,\n 'vacuum_min_vector_number': 1000,\n 'default_segment_number': 0,\n 'max_segment_size': None,\n 'memmap_threshold': None,\n 'indexing_threshold': 20000,\n 'flush_interval_sec': 5,\n 'max_optimization_threads': None},\n 'wal_config': {'wal_capacity_mb': 32, 'wal_segments_ahead': 0},\n 'quantization_config': {'binary': {'always_ram': True}}},\n 'payload_schema': {}}"
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"collection_info = client.get_collection(collection_name=f\"{collection_name}\")\n",
|
||||
"collection_info.dict()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Oversampling vs Recall\n",
|
||||
"\n",
|
||||
"### Preparing a query dataset\n",
|
||||
"\n",
|
||||
"For the purpose of this illustration, we'll take a few vectors which we know are already in the index and query them. We should get the same vectors back as results from the Qdrant index. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:34:13.453626Z",
|
||||
"start_time": "2024-04-01T16:34:13.391567Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "[89391,\n 79659,\n 12006,\n 80978,\n 87219,\n 97885,\n 83155,\n 67504,\n 4645,\n 82711,\n 48395,\n 57375,\n 69208,\n 14136,\n 89515,\n 59880,\n 78730,\n 36952,\n 49620,\n 96486,\n 55473,\n 58179,\n 18926,\n 6489,\n 11931,\n 54146,\n 9850,\n 71259,\n 37825,\n 47331,\n 84964,\n 92399,\n 56669,\n 77042,\n 73744,\n 47993,\n 83780,\n 92429,\n 75114,\n 4463,\n 69030,\n 81185,\n 27950,\n 66217,\n 54652,\n 8260,\n 1151,\n 993,\n 85954,\n 66863,\n 47303,\n 8992,\n 92688,\n 76030,\n 29472,\n 3077,\n 42454,\n 46120,\n 69140,\n 20877,\n 2844,\n 95423,\n 1770,\n 28568,\n 96448,\n 94227,\n 40837,\n 91684,\n 29785,\n 66936,\n 85121,\n 39546,\n 81910,\n 5514,\n 37068,\n 35731,\n 93990,\n 26685,\n 63076,\n 18762,\n 27922,\n 34916,\n 80976,\n 83189,\n 6328,\n 57508,\n 58860,\n 13758,\n 72976,\n 85030,\n 332,\n 34963,\n 85009,\n 31344,\n 11560,\n 58108,\n 85163,\n 17064,\n 44712,\n 45962]"
|
||||
},
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"query_indices = random.sample(range(len(dataset)), 100)\n",
|
||||
"query_dataset = dataset[query_indices]\n",
|
||||
"query_indices"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:34:13.453928Z",
|
||||
"start_time": "2024-04-01T16:34:13.452405Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"## Add Gaussian noise to any vector\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def add_noise(vector, noise=0.05):\n",
|
||||
" return vector + noise * np.random.randn(*vector.shape)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:34:13.457839Z",
|
||||
"start_time": "2024-04-01T16:34:13.455431Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def correct(results, text):\n",
|
||||
" return text in [x.payload[\"text\"] for x in results]\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def count_correct(query_dataset, limit=1, oversampling=1, rescore=False):\n",
|
||||
" correct_results = 0\n",
|
||||
" for query_vector, text in zip(query_dataset[\"openai\"], query_dataset[\"text\"]):\n",
|
||||
" results = client.search(\n",
|
||||
" collection_name=collection_name,\n",
|
||||
" query_vector=add_noise(np.array(query_vector)),\n",
|
||||
" limit=limit,\n",
|
||||
" search_params=models.SearchParams(\n",
|
||||
" quantization=models.QuantizationSearchParams(\n",
|
||||
" rescore=rescore,\n",
|
||||
" oversampling=oversampling,\n",
|
||||
" )\n",
|
||||
" ),\n",
|
||||
" )\n",
|
||||
" correct_results += correct(results, text)\n",
|
||||
" return correct_results"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:40:48.347002Z",
|
||||
"start_time": "2024-04-01T16:40:42.228551Z"
|
||||
},
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"limit_grid = [1, 3, 10, 20, 50]\n",
|
||||
"oversampling_grid = [1.0, 3.0, 5.0]\n",
|
||||
"rescore_grid = [False, True]\n",
|
||||
"results = []\n",
|
||||
"\n",
|
||||
"for limit in limit_grid:\n",
|
||||
" for oversampling in oversampling_grid:\n",
|
||||
" for rescore in rescore_grid:\n",
|
||||
" start = time.perf_counter()\n",
|
||||
" correct_results = count_correct(\n",
|
||||
" query_dataset, limit=limit, oversampling=oversampling, rescore=rescore\n",
|
||||
" )\n",
|
||||
" end = time.perf_counter()\n",
|
||||
" results.append(\n",
|
||||
" {\n",
|
||||
" \"limit\": limit,\n",
|
||||
" \"oversampling\": oversampling,\n",
|
||||
" \"candidates\": int(oversampling * limit),\n",
|
||||
" \"rescore\": rescore,\n",
|
||||
" \"accuracy\": correct_results / 100,\n",
|
||||
" \"total queries\": len(query_dataset[\"text\"]),\n",
|
||||
" \"time\": end - start,\n",
|
||||
" }\n",
|
||||
" )"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 22,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-04-01T16:41:55.445405Z",
|
||||
"start_time": "2024-04-01T16:41:55.442687Z"
|
||||
}
|
||||
},
|
||||
"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>candidates</th>\n <th>rescore</th>\n <th>accuracy</th>\n <th>time</th>\n </tr>\n </thead>\n <tbody>\n <tr>\n <th>0</th>\n <td>1</td>\n <td>False</td>\n <td>0.90</td>\n <td>0.221826</td>\n </tr>\n <tr>\n <th>1</th>\n <td>1</td>\n <td>True</td>\n <td>0.91</td>\n <td>0.134167</td>\n </tr>\n <tr>\n <th>2</th>\n <td>3</td>\n <td>False</td>\n <td>0.88</td>\n <td>0.115299</td>\n </tr>\n <tr>\n <th>3</th>\n <td>3</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.209320</td>\n </tr>\n <tr>\n <th>4</th>\n <td>5</td>\n <td>False</td>\n <td>0.84</td>\n <td>0.154485</td>\n </tr>\n <tr>\n <th>5</th>\n <td>5</td>\n <td>True</td>\n <td>0.91</td>\n <td>0.124424</td>\n </tr>\n <tr>\n <th>6</th>\n <td>3</td>\n <td>False</td>\n <td>0.99</td>\n <td>0.121695</td>\n </tr>\n <tr>\n <th>7</th>\n <td>3</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.123257</td>\n </tr>\n <tr>\n <th>8</th>\n <td>9</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.119629</td>\n </tr>\n <tr>\n <th>9</th>\n <td>9</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.119372</td>\n </tr>\n <tr>\n <th>10</th>\n <td>15</td>\n <td>False</td>\n <td>0.90</td>\n <td>0.121621</td>\n </tr>\n <tr>\n <th>11</th>\n <td>15</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.125466</td>\n </tr>\n <tr>\n <th>12</th>\n <td>10</td>\n <td>False</td>\n <td>0.93</td>\n <td>0.135910</td>\n </tr>\n <tr>\n <th>13</th>\n <td>10</td>\n <td>True</td>\n <td>0.95</td>\n <td>0.138135</td>\n </tr>\n <tr>\n <th>14</th>\n <td>30</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.177928</td>\n </tr>\n <tr>\n <th>15</th>\n <td>30</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.254588</td>\n </tr>\n <tr>\n <th>16</th>\n <td>50</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.268659</td>\n </tr>\n <tr>\n <th>17</th>\n <td>50</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.269792</td>\n </tr>\n <tr>\n <th>18</th>\n <td>20</td>\n <td>False</td>\n <td>0.96</td>\n <td>0.249941</td>\n </tr>\n <tr>\n <th>19</th>\n <td>20</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.247138</td>\n </tr>\n <tr>\n <th>20</th>\n <td>60</td>\n <td>False</td>\n <td>0.97</td>\n <td>0.251301</td>\n </tr>\n <tr>\n <th>21</th>\n <td>60</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.256504</td>\n </tr>\n <tr>\n <th>22</th>\n <td>100</td>\n <td>False</td>\n <td>0.98</td>\n <td>0.270049</td>\n </tr>\n <tr>\n <th>23</th>\n <td>100</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.248972</td>\n </tr>\n <tr>\n <th>24</th>\n <td>50</td>\n <td>False</td>\n <td>0.97</td>\n <td>0.306356</td>\n </tr>\n <tr>\n <th>25</th>\n <td>50</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.257544</td>\n </tr>\n <tr>\n <th>26</th>\n <td>150</td>\n <td>False</td>\n <td>0.98</td>\n <td>0.238811</td>\n </tr>\n <tr>\n <th>27</th>\n <td>150</td>\n <td>True</td>\n <td>0.99</td>\n <td>0.263939</td>\n </tr>\n <tr>\n <th>28</th>\n <td>250</td>\n <td>False</td>\n <td>0.99</td>\n <td>0.256558</td>\n </tr>\n <tr>\n <th>29</th>\n <td>250</td>\n <td>True</td>\n <td>1.00</td>\n <td>0.335823</td>\n </tr>\n </tbody>\n</table>\n</div>",
|
||||
"text/plain": " candidates rescore accuracy time\n0 1 False 0.90 0.221826\n1 1 True 0.91 0.134167\n2 3 False 0.88 0.115299\n3 3 True 0.97 0.209320\n4 5 False 0.84 0.154485\n5 5 True 0.91 0.124424\n6 3 False 0.99 0.121695\n7 3 True 0.96 0.123257\n8 9 False 0.94 0.119629\n9 9 True 0.98 0.119372\n10 15 False 0.90 0.121621\n11 15 True 0.97 0.125466\n12 10 False 0.93 0.135910\n13 10 True 0.95 0.138135\n14 30 False 0.94 0.177928\n15 30 True 0.98 0.254588\n16 50 False 0.94 0.268659\n17 50 True 0.96 0.269792\n18 20 False 0.96 0.249941\n19 20 True 0.96 0.247138\n20 60 False 0.97 0.251301\n21 60 True 0.98 0.256504\n22 100 False 0.98 0.270049\n23 100 True 0.97 0.248972\n24 50 False 0.97 0.306356\n25 50 True 0.98 0.257544\n26 150 False 0.98 0.238811\n27 150 True 0.99 0.263939\n28 250 False 0.99 0.256558\n29 250 True 1.00 0.335823"
|
||||
},
|
||||
"execution_count": 22,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"df = pd.DataFrame(results)\n",
|
||||
"df[[\"candidates\", \"rescore\", \"accuracy\", \"time\"]]\n",
|
||||
"# df.to_csv(\"candidates-rescore-time.csv\", index=False)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -1,371 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"import torch\n",
|
||||
"from transformers import AutoModelForMaskedLM, AutoTokenizer"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Running the model with Transformers and Torch"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"sentences = [\n",
|
||||
" \"Hello World\",\n",
|
||||
" \"Built by Nirant Kasliwal\",\n",
|
||||
"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## PyTorch Code from the [SPLADERunner](https://github.com/PrithivirajDamodaran/SPLADERunner) library"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"hf_token = \"<your_hf_token_here>\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Output Logits shape: torch.Size([2, 10, 30522])\n",
|
||||
"Output Attention mask shape: torch.Size([2, 10])\n",
|
||||
"Sparse Vector shape: torch.Size([2, 30522])\n",
|
||||
"SPLADE BOW rep for sentence:\tBuilt by Nirant Kasliwal\n",
|
||||
"[('##rant', 2.02), ('built', 1.94), ('##wal', 1.79), ('##sl', 1.69), ('build', 1.57), ('ka', 1.4), ('ni', 1.26), ('made', 0.93), ('architect', 0.76), ('was', 0.69), ('who', 0.61), ('his', 0.5), ('wrote', 0.47), ('india', 0.45), ('company', 0.41), ('##i', 0.41), ('he', 0.37), ('manufacturer', 0.36), ('by', 0.35), ('engineer', 0.33), ('architecture', 0.33), ('ko', 0.23), ('him', 0.22), ('invented', 0.19), ('said', 0.14), ('k', 0.11), ('man', 0.11), ('statue', 0.11), ('bomb', 0.1), ('##wa', 0.1), ('builder', 0.09), ('.', 0.07), ('started', 0.06), (',', 0.04), ('ku', 0.03)]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Download the model and tokenizer\n",
|
||||
"device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"prithivida/Splade_PP_en_v1\", token=hf_token)\n",
|
||||
"reverse_voc = {v: k for k, v in tokenizer.vocab.items()}\n",
|
||||
"model = AutoModelForMaskedLM.from_pretrained(\"prithivida/Splade_PP_en_v1\", token=hf_token)\n",
|
||||
"model.to(device)\n",
|
||||
"\n",
|
||||
"# Tokenize the input\n",
|
||||
"inputs = tokenizer(sentences, return_tensors=\"pt\", padding=True, truncation=True, max_length=512)\n",
|
||||
"inputs = {key: val.to(device) for key, val in inputs.items()}\n",
|
||||
"input_ids = inputs[\"input_ids\"]\n",
|
||||
"attention_mask = inputs[\"attention_mask\"]\n",
|
||||
"token_type_ids = inputs[\"token_type_ids\"]\n",
|
||||
"\n",
|
||||
"# Run model and prepare sparse vector\n",
|
||||
"outputs = model(**inputs)\n",
|
||||
"logits = outputs.logits\n",
|
||||
"print(\"Output Logits shape: \", logits.shape)\n",
|
||||
"print(\"Output Attention mask shape: \", attention_mask.shape)\n",
|
||||
"relu_log = torch.log(1 + torch.relu(logits))\n",
|
||||
"weighted_log = relu_log * attention_mask.unsqueeze(-1)\n",
|
||||
"max_val, _ = torch.max(weighted_log, dim=1)\n",
|
||||
"vector = max_val.squeeze()\n",
|
||||
"print(\"Sparse Vector shape: \", vector.shape)\n",
|
||||
"# print(\"Number of Actual Dimensions: \", len(cols))\n",
|
||||
"cols = [vec.nonzero().squeeze().cpu().tolist() for vec in vector]\n",
|
||||
"weights = [vec[col].cpu().tolist() for vec, col in zip(vector, cols)]\n",
|
||||
"\n",
|
||||
"idx = 1\n",
|
||||
"cols, weights = cols[idx], weights[idx]\n",
|
||||
"# Print the BOW representation\n",
|
||||
"d = {k: v for k, v in zip(cols, weights)}\n",
|
||||
"sorted_d = {k: v for k, v in sorted(d.items(), key=lambda item: item[1], reverse=True)}\n",
|
||||
"bow_rep = []\n",
|
||||
"for k, v in sorted_d.items():\n",
|
||||
" bow_rep.append((reverse_voc[k], round(v, 2)))\n",
|
||||
"print(f\"SPLADE BOW rep for sentence:\\t{sentences[idx]}\\n{bow_rep}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Export with output_attentions and logits"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Exporting model to models/nirantk_SPLADE_PP_en_v1\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"('models/nirantk_SPLADE_PP_en_v1/tokenizer_config.json',\n",
|
||||
" 'models/nirantk_SPLADE_PP_en_v1/special_tokens_map.json',\n",
|
||||
" 'models/nirantk_SPLADE_PP_en_v1/vocab.txt',\n",
|
||||
" 'models/nirantk_SPLADE_PP_en_v1/added_tokens.json',\n",
|
||||
" 'models/nirantk_SPLADE_PP_en_v1/tokenizer.json')"
|
||||
]
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from transformers import AutoTokenizer\n",
|
||||
"\n",
|
||||
"model_id = \"nirantk/SPLADE_PP_en_v1\"\n",
|
||||
"output_dir = f\"models/{model_id.replace('/', '_')}\"\n",
|
||||
"model_kwargs = {\"output_attentions\": True, \"return_dict\": True}\n",
|
||||
"\n",
|
||||
"print(f\"Exporting model to {output_dir}\")\n",
|
||||
"tokenizer.save_pretrained(output_dir)\n",
|
||||
"# main_export(\n",
|
||||
"# model_id,\n",
|
||||
"# output=output_dir,\n",
|
||||
"# no_post_process=True,\n",
|
||||
"# model_kwargs=model_kwargs,\n",
|
||||
"# token=hf_token,\n",
|
||||
"# )"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Running the model with ONNX"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from optimum.onnxruntime import ORTModelForMaskedLM\n",
|
||||
"\n",
|
||||
"model = ORTModelForMaskedLM.from_pretrained(\"nirantk/SPLADE_PP_en_v1\")\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"nirantk/SPLADE_PP_en_v1\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"inputs = tokenizer(sentences, return_tensors=\"pt\", padding=True, truncation=True, max_length=512)\n",
|
||||
"inputs = {key: val.to(device) for key, val in inputs.items()}\n",
|
||||
"input_ids = inputs[\"input_ids\"]\n",
|
||||
"attention_mask = inputs[\"attention_mask\"]\n",
|
||||
"token_type_ids = inputs[\"token_type_ids\"]\n",
|
||||
"\n",
|
||||
"onnx_input = {\n",
|
||||
" \"input_ids\": input_ids.cpu().numpy(),\n",
|
||||
" \"attention_mask\": attention_mask.cpu().numpy(),\n",
|
||||
" \"token_type_ids\": token_type_ids.cpu().numpy(),\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"logits = model(**onnx_input).logits"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"(2, 10, 30522)"
|
||||
]
|
||||
},
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"logits.shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Output Logits shape: (2, 10, 30522)\n",
|
||||
"Sparse Vector shape: (2, 30522)\n",
|
||||
"SPLADE BOW rep for sentence:\tBuilt by Nirant Kasliwal\n",
|
||||
"[('##rant', 2.02), ('built', 1.94), ('##wal', 1.79), ('##sl', 1.69), ('build', 1.57), ('ka', 1.4), ('ni', 1.26), ('made', 0.93), ('architect', 0.76), ('was', 0.69), ('who', 0.61), ('his', 0.5), ('wrote', 0.47), ('india', 0.45), ('company', 0.41), ('##i', 0.41), ('he', 0.37), ('manufacturer', 0.36), ('by', 0.35), ('engineer', 0.33), ('architecture', 0.33), ('ko', 0.23), ('him', 0.22), ('invented', 0.19), ('said', 0.14), ('k', 0.11), ('man', 0.11), ('statue', 0.11), ('bomb', 0.1), ('##wa', 0.1), ('builder', 0.09), ('.', 0.07), ('started', 0.06), (',', 0.04), ('ku', 0.03)]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(\"Output Logits shape: \", logits.shape)\n",
|
||||
"\n",
|
||||
"relu_log = np.log(1 + np.maximum(logits, 0))\n",
|
||||
"\n",
|
||||
"# Equivalent to relu_log * attention_mask.unsqueeze(-1)\n",
|
||||
"# For NumPy, you might need to explicitly expand dimensions if 'attention_mask' is not already 2D\n",
|
||||
"weighted_log = relu_log * np.expand_dims(attention_mask, axis=-1)\n",
|
||||
"\n",
|
||||
"# Equivalent to torch.max(weighted_log, dim=1)\n",
|
||||
"# NumPy's max function returns only the max values, not the indices, so we don't need to unpack two values\n",
|
||||
"max_val = np.max(weighted_log, axis=1)\n",
|
||||
"\n",
|
||||
"# Equivalent to max_val.squeeze()\n",
|
||||
"# This step may be unnecessary in NumPy if max_val doesn't have unnecessary dimensions\n",
|
||||
"vector = np.squeeze(max_val)\n",
|
||||
"print(\"Sparse Vector shape: \", vector.shape)\n",
|
||||
"\n",
|
||||
"# print(vector[0].nonzero())\n",
|
||||
"\n",
|
||||
"cols = [vec.nonzero()[0].squeeze().tolist() for vec in vector]\n",
|
||||
"weights = [vec[col].tolist() for vec, col in zip(vector, cols)]\n",
|
||||
"\n",
|
||||
"idx = 1\n",
|
||||
"cols, weights = cols[idx], weights[idx]\n",
|
||||
"# Print the BOW representation\n",
|
||||
"d = {k: v for k, v in zip(cols, weights)}\n",
|
||||
"sorted_d = {k: v for k, v in sorted(d.items(), key=lambda item: item[1], reverse=True)}\n",
|
||||
"bow_rep = []\n",
|
||||
"for k, v in sorted_d.items():\n",
|
||||
" bow_rep.append((reverse_voc[k], round(v, 2)))\n",
|
||||
"print(f\"SPLADE BOW rep for sentence:\\t{sentences[idx]}\\n{bow_rep}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"35"
|
||||
]
|
||||
},
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"len(cols)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[1010,\n",
|
||||
" 1012,\n",
|
||||
" 1047,\n",
|
||||
" 2001,\n",
|
||||
" 2002,\n",
|
||||
" 2010,\n",
|
||||
" 2011,\n",
|
||||
" 2032,\n",
|
||||
" 2040,\n",
|
||||
" 2056,\n",
|
||||
" 2072,\n",
|
||||
" 2081,\n",
|
||||
" 2158,\n",
|
||||
" 2194,\n",
|
||||
" 2318,\n",
|
||||
" 2328,\n",
|
||||
" 2626,\n",
|
||||
" 2634,\n",
|
||||
" 3857,\n",
|
||||
" 3992,\n",
|
||||
" 4213,\n",
|
||||
" 4294,\n",
|
||||
" 4944,\n",
|
||||
" 5968,\n",
|
||||
" 6231,\n",
|
||||
" 7751,\n",
|
||||
" 8826,\n",
|
||||
" 9152,\n",
|
||||
" 10556,\n",
|
||||
" 12508,\n",
|
||||
" 12849,\n",
|
||||
" 13476,\n",
|
||||
" 13970,\n",
|
||||
" 14540,\n",
|
||||
" 17884]"
|
||||
]
|
||||
},
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"cols"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "fst",
|
||||
"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"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -1,122 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "4bdb2a91-fa2a-4cee-ad5a-176cc957394d",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-23T12:15:28.171586Z",
|
||||
"start_time": "2024-05-23T12:15:28.076314Z"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"ename": "ModuleNotFoundError",
|
||||
"evalue": "No module named 'torch'",
|
||||
"output_type": "error",
|
||||
"traceback": [
|
||||
"\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
|
||||
"\u001B[0;31mModuleNotFoundError\u001B[0m Traceback (most recent call last)",
|
||||
"Cell \u001B[0;32mIn[1], line 1\u001B[0m\n\u001B[0;32m----> 1\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\n\u001B[1;32m 2\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01monnx\u001B[39;00m\n\u001B[1;32m 3\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorchvision\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mmodels\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m \u001B[38;5;21;01mmodels\u001B[39;00m\n",
|
||||
"\u001B[0;31mModuleNotFoundError\u001B[0m: No module named 'torch'"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"import torch.onnx\n",
|
||||
"import torchvision.models as models\n",
|
||||
"import torchvision.transforms as transforms\n",
|
||||
"from PIL import Image\n",
|
||||
"import numpy as np\n",
|
||||
"from tests.config import TEST_MISC_DIR\n",
|
||||
"\n",
|
||||
"# Load pre-trained ResNet-50 model\n",
|
||||
"resnet = models.resnet50(pretrained=True)\n",
|
||||
"resnet = torch.nn.Sequential(*(list(resnet.children())[:-1])) # Remove the last fully connected layer\n",
|
||||
"resnet.eval()\n",
|
||||
"\n",
|
||||
"# Define preprocessing transform\n",
|
||||
"preprocess = transforms.Compose([\n",
|
||||
" transforms.Resize(256),\n",
|
||||
" transforms.CenterCrop(224),\n",
|
||||
" transforms.ToTensor(),\n",
|
||||
" transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
|
||||
"])\n",
|
||||
"\n",
|
||||
"# Load and preprocess the image\n",
|
||||
"def preprocess_image(image_path):\n",
|
||||
" input_image = Image.open(image_path)\n",
|
||||
" input_tensor = preprocess(input_image)\n",
|
||||
" input_batch = input_tensor.unsqueeze(0) # Add batch dimension\n",
|
||||
" return input_batch\n",
|
||||
"\n",
|
||||
"# Example input for exporting\n",
|
||||
"input_image = preprocess_image('example.jpg')\n",
|
||||
"\n",
|
||||
"# Export the model to ONNX with dynamic axes\n",
|
||||
"torch.onnx.export(\n",
|
||||
" resnet, \n",
|
||||
" input_image, \n",
|
||||
" \"model.onnx\", \n",
|
||||
" export_params=True, \n",
|
||||
" opset_version=9, \n",
|
||||
" input_names=['input'], \n",
|
||||
" output_names=['output'],\n",
|
||||
" dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Load ONNX model\n",
|
||||
"import onnx\n",
|
||||
"import onnxruntime as ort\n",
|
||||
"\n",
|
||||
"onnx_model = onnx.load(\"model.onnx\")\n",
|
||||
"ort_session = ort.InferenceSession(\"model.onnx\")\n",
|
||||
"\n",
|
||||
"# Run inference and extract feature vectors\n",
|
||||
"def extract_feature_vectors(image_paths):\n",
|
||||
" input_images = [preprocess_image(image_path) for image_path in image_paths]\n",
|
||||
" input_batch = torch.cat(input_images, dim=0) # Combine images into a single batch\n",
|
||||
" ort_inputs = {ort_session.get_inputs()[0].name: input_batch.numpy()}\n",
|
||||
" ort_outs = ort_session.run(None, ort_inputs)\n",
|
||||
" return ort_outs[0]\n",
|
||||
"\n",
|
||||
"# Example usage\n",
|
||||
"images = [TEST_MISC_DIR / \"image.jpeg\", str(TEST_MISC_DIR / \"small_image.jpeg\")] # Replace with your image paths\n",
|
||||
"feature_vectors = extract_feature_vectors(images)\n",
|
||||
"print(\"Feature vector shape:\", feature_vectors.shape)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"outputs": [],
|
||||
"source": [],
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"id": "baa650c4cb3e0e6d"
|
||||
}
|
||||
],
|
||||
"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.12.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -1,15 +0,0 @@
|
||||
from optimum.exporters.onnx import main_export
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
model_id = "sentence-transformers/paraphrase-MiniLM-L6-v2"
|
||||
output_dir = f"models/{model_id.replace('/', '_')}"
|
||||
model_kwargs = {"output_attentions": True, "return_dict": True}
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
|
||||
# export if the output model does not exist
|
||||
# try:
|
||||
# sess = onnxruntime.InferenceSession(f"{output_dir}/model.onnx")
|
||||
# print("Model already exported")
|
||||
# except FileNotFoundError:
|
||||
print(f"Exporting model to {output_dir}")
|
||||
main_export(model_id, output=output_dir, no_post_process=True, model_kwargs=model_kwargs)
|
||||
@@ -0,0 +1,473 @@
|
||||
"""
|
||||
This script is used to convert HuggingFace models to ONNX format and optionally quantize the model using dynamic quantization.
|
||||
This is courtesy of Joshua aka @Xenova
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Set
|
||||
|
||||
import onnx
|
||||
from onnxruntime.quantization import QuantType, quantize_dynamic
|
||||
from optimum.exporters.onnx import export_models, main_export
|
||||
from optimum.exporters.tasks import TasksManager
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoConfig, AutoTokenizer, HfArgumentParser
|
||||
|
||||
DEFAULT_QUANTIZE_PARAMS = {
|
||||
"per_channel": True,
|
||||
"reduce_range": True,
|
||||
}
|
||||
|
||||
MODEL_SPECIFIC_QUANTIZE_PARAMS = {
|
||||
# Decoder-only models
|
||||
"codegen": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt2": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt_bigcode": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gptj": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt-neo": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt-neox": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"mpt": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"bloom": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"llama": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"opt": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"mistral": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"falcon": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"phi": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"qwen2": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
# Encoder-decoder models
|
||||
"whisper": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"vision-encoder-decoder": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
# Encoder-only models
|
||||
"owlv2": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
}
|
||||
|
||||
MODELS_WITHOUT_TOKENIZERS = [
|
||||
"wav2vec2",
|
||||
"wav2vec2-bert",
|
||||
"wavlm",
|
||||
"hubert",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversionArguments:
|
||||
"""
|
||||
Arguments used for converting HuggingFace models to onnx.
|
||||
"""
|
||||
|
||||
model_id: str = field(metadata={"help": "Model identifier"})
|
||||
tokenizer_id: str = field(default=None, metadata={"help": "Tokenizer identifier (if different to `model_id`)"})
|
||||
quantize: bool = field(default=False, metadata={"help": "Whether to quantize the model."})
|
||||
output_parent_dir: str = field(
|
||||
default="./models/", metadata={"help": "Path where the converted model will be saved to."}
|
||||
)
|
||||
|
||||
task: Optional[str] = field(
|
||||
default="auto",
|
||||
metadata={
|
||||
"help": (
|
||||
"The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
|
||||
f" {str(TasksManager.get_all_tasks())}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
opset: int = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": (
|
||||
"If specified, ONNX opset version to export the model with. Otherwise, the default opset will be used."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
device: str = field(default="cpu", metadata={"help": "The device to use to do the export."})
|
||||
skip_validation: bool = field(default=False, metadata={"help": "Whether to skip validation of the converted model"})
|
||||
|
||||
per_channel: bool = field(default=None, metadata={"help": "Whether to quantize weights per channel"})
|
||||
reduce_range: bool = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "Whether to quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode"
|
||||
},
|
||||
)
|
||||
|
||||
output_attentions: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "Whether to output attentions from the model. NOTE: This is only supported for whisper models right now."
|
||||
},
|
||||
)
|
||||
|
||||
split_modalities: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "Whether to split multimodal models. NOTE: This is only supported for CLIP models right now."
|
||||
},
|
||||
)
|
||||
|
||||
trust_remote_code: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories"
|
||||
"you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository."
|
||||
},
|
||||
)
|
||||
|
||||
custom_onnx_configs: str = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "Experimental usage: override the default ONNX config used for the given model. This argument may be useful for advanced users "
|
||||
"that desire a finer-grained control on the export."
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def get_operators(model: onnx.ModelProto) -> Set[str]:
|
||||
operators = set()
|
||||
|
||||
def traverse_graph(graph):
|
||||
for node in graph.node:
|
||||
operators.add(node.op_type)
|
||||
for attr in node.attribute:
|
||||
if attr.type == onnx.AttributeProto.GRAPH:
|
||||
subgraph = attr.g
|
||||
traverse_graph(subgraph)
|
||||
|
||||
traverse_graph(model.graph)
|
||||
return operators
|
||||
|
||||
|
||||
def quantize(model_names_or_paths, **quantize_kwargs):
|
||||
"""
|
||||
Quantize the weights of the model from float32 to int8 to allow very efficient inference on modern CPU
|
||||
|
||||
Uses unsigned ints for activation values, signed ints for weights, per
|
||||
https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection
|
||||
it is faster on most CPU architectures
|
||||
Args:
|
||||
onnx_model_path: Path to location the exported ONNX model is stored
|
||||
Returns: The Path generated for the quantized
|
||||
"""
|
||||
|
||||
quantize_config = dict(**quantize_kwargs, per_model_config={})
|
||||
|
||||
for model in tqdm(model_names_or_paths, desc="Quantizing"):
|
||||
directory_path = os.path.dirname(model)
|
||||
file_name_without_extension = os.path.splitext(os.path.basename(model))[0]
|
||||
|
||||
# NOTE:
|
||||
# As of 2023/04/20, the current latest version of onnxruntime-web is 1.14.0, and does not support INT8 weights for Conv layers.
|
||||
# For this reason, we choose model weight types to ensure compatibility with onnxruntime-web.
|
||||
#
|
||||
# As per docs, signed weight type (QInt8) is faster on most CPUs, so, we use that unless the model contains a Conv layer.
|
||||
# For more information, see:
|
||||
# - https://github.com/microsoft/onnxruntime/issues/3130#issuecomment-1105200621
|
||||
# - https://github.com/microsoft/onnxruntime/issues/2339
|
||||
|
||||
loaded_model = onnx.load_model(model)
|
||||
op_types = get_operators(loaded_model)
|
||||
weight_type = QuantType.QUInt8 if "Conv" in op_types else QuantType.QInt8
|
||||
|
||||
quantize_dynamic(
|
||||
model_input=model,
|
||||
model_output=os.path.join(directory_path, f"{file_name_without_extension}_quantized.onnx"),
|
||||
weight_type=weight_type,
|
||||
# TODO allow user to specify these
|
||||
# op_types_to_quantize=['MatMul', 'Add', 'Conv'],
|
||||
extra_options=dict(EnableSubgraph=True),
|
||||
**quantize_kwargs,
|
||||
)
|
||||
|
||||
quantize_config["per_model_config"][file_name_without_extension] = dict(
|
||||
op_types=list(op_types),
|
||||
weight_type=str(weight_type),
|
||||
)
|
||||
|
||||
# Save quantization config
|
||||
with open(os.path.join(directory_path, "quantize_config.json"), "w") as fp:
|
||||
json.dump(quantize_config, fp, indent=4)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((ConversionArguments,))
|
||||
(conv_args,) = parser.parse_args_into_dataclasses()
|
||||
|
||||
model_id = conv_args.model_id
|
||||
tokenizer_id = conv_args.tokenizer_id or model_id
|
||||
|
||||
output_model_folder = os.path.join(conv_args.output_parent_dir, model_id)
|
||||
|
||||
# Create output folder
|
||||
os.makedirs(output_model_folder, exist_ok=True)
|
||||
|
||||
from_pretrained_kwargs = dict(
|
||||
trust_remote_code=conv_args.trust_remote_code,
|
||||
)
|
||||
|
||||
# Saving the model config
|
||||
config = AutoConfig.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
custom_kwargs = {}
|
||||
if conv_args.custom_onnx_configs is not None:
|
||||
if conv_args.task == "auto":
|
||||
raise Exception("`--task` must be set when exporting with `--custom_onnx_configs`")
|
||||
custom_onnx_configs = json.loads(conv_args.custom_onnx_configs)
|
||||
|
||||
for key in custom_onnx_configs:
|
||||
onnx_configs = TasksManager._SUPPORTED_MODEL_TYPE[custom_onnx_configs[key]]["onnx"]
|
||||
mapping = onnx_configs[conv_args.task]
|
||||
custom_onnx_configs[key] = mapping.func(config, **mapping.keywords)
|
||||
|
||||
custom_kwargs["custom_onnx_configs"] = custom_onnx_configs
|
||||
|
||||
tokenizer = None
|
||||
try:
|
||||
# Load tokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, **from_pretrained_kwargs)
|
||||
|
||||
# To avoid inserting all chat templates into tokenizers.js, we save the chat template
|
||||
# to the tokenizer_config.json file, and load it when the tokenizer is loaded.
|
||||
if getattr(tokenizer, "chat_template", None) is None and getattr(tokenizer, "use_default_system_prompt", False):
|
||||
# No chat template specified, and we use the default
|
||||
setattr(tokenizer, "chat_template", tokenizer.default_chat_template)
|
||||
|
||||
except KeyError:
|
||||
pass # No Tokenizer
|
||||
|
||||
except Exception as e:
|
||||
if config.model_type not in MODELS_WITHOUT_TOKENIZERS:
|
||||
raise e
|
||||
|
||||
core_export_kwargs = dict(
|
||||
opset=conv_args.opset,
|
||||
device=conv_args.device,
|
||||
trust_remote_code=conv_args.trust_remote_code,
|
||||
**custom_kwargs,
|
||||
)
|
||||
|
||||
export_kwargs = dict(
|
||||
model_name_or_path=model_id,
|
||||
output=output_model_folder,
|
||||
task=conv_args.task,
|
||||
do_validation=not conv_args.skip_validation,
|
||||
library_name="transformers",
|
||||
**core_export_kwargs,
|
||||
)
|
||||
|
||||
# Handle special cases
|
||||
if config.model_type == "marian":
|
||||
from .extra.marian import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(model_id, tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type == "esm":
|
||||
from .extra.esm import generate_fast_tokenizer
|
||||
|
||||
fast_tokenizer = generate_fast_tokenizer(tokenizer)
|
||||
fast_tokenizer.save(os.path.join(output_model_folder, "tokenizer.json"))
|
||||
|
||||
elif config.model_type == "whisper":
|
||||
if conv_args.output_attentions:
|
||||
from .extra.whisper import get_main_export_kwargs
|
||||
|
||||
export_kwargs.update(**get_main_export_kwargs(config, "automatic-speech-recognition"))
|
||||
|
||||
elif config.model_type in ("wav2vec2", "wav2vec2-bert", "hubert"):
|
||||
if tokenizer is not None:
|
||||
from .extra.wav2vec2 import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type == "vits":
|
||||
if tokenizer is not None:
|
||||
from .extra.vits import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type == "speecht5":
|
||||
# TODO allow user to specify vocoder path
|
||||
export_kwargs["model_kwargs"] = {"vocoder": "microsoft/speecht5_hifigan"}
|
||||
|
||||
if tokenizer is not None:
|
||||
from .extra.speecht5 import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type in ("owlvit", "owlv2"):
|
||||
# Override default batch size to 1, needed because non-maximum suppression is performed for exporting.
|
||||
# For more information, see https://github.com/huggingface/optimum/blob/e3b7efb1257c011db907ef40ab340e795cc5684c/optimum/exporters/onnx/model_configs.py#L1028-L1032
|
||||
export_kwargs["batch_size"] = 1
|
||||
|
||||
else:
|
||||
pass # TODO
|
||||
|
||||
# Step 1. convert huggingface model to onnx
|
||||
if not conv_args.split_modalities:
|
||||
main_export(**export_kwargs)
|
||||
else:
|
||||
custom_export_kwargs = dict(
|
||||
output_dir=output_model_folder,
|
||||
**core_export_kwargs,
|
||||
)
|
||||
|
||||
if config.model_type == "clip":
|
||||
# Handle special case for exporting text and vision models separately
|
||||
from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection
|
||||
|
||||
from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig
|
||||
|
||||
text_model = CLIPTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
export_models(
|
||||
models_and_onnx_configs={
|
||||
"text_model": (text_model, CLIPTextModelWithProjectionOnnxConfig(text_model.config)),
|
||||
"vision_model": (vision_model, CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)),
|
||||
},
|
||||
**custom_export_kwargs,
|
||||
)
|
||||
|
||||
elif config.model_type == "siglip":
|
||||
# Handle special case for exporting text and vision models separately
|
||||
from transformers.models.siglip import SiglipTextModel, SiglipVisionModel
|
||||
|
||||
from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig
|
||||
|
||||
text_model = SiglipTextModel.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
vision_model = SiglipVisionModel.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
export_models(
|
||||
models_and_onnx_configs={
|
||||
"text_model": (text_model, SiglipTextModelOnnxConfig(text_model.config)),
|
||||
"vision_model": (vision_model, SiglipVisionModelOnnxConfig(vision_model.config)),
|
||||
},
|
||||
**custom_export_kwargs,
|
||||
)
|
||||
|
||||
# TODO: Enable once https://github.com/huggingface/optimum/pull/1552 is merged
|
||||
# elif config.model_type == 'clap':
|
||||
# # Handle special case for exporting text and audio models separately
|
||||
# from .extra.clap import ClapTextModelWithProjectionOnnxConfig, ClapAudioModelWithProjectionOnnxConfig
|
||||
# from transformers.models.clap import ClapTextModelWithProjection, ClapAudioModelWithProjection
|
||||
|
||||
# text_model = ClapTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
# audio_model = ClapAudioModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
# export_models(
|
||||
# models_and_onnx_configs={
|
||||
# "text_model": (text_model, ClapTextModelWithProjectionOnnxConfig(text_model.config)),
|
||||
# "audio_model": (audio_model, ClapAudioModelWithProjectionOnnxConfig(audio_model.config)),
|
||||
# },
|
||||
# **custom_export_kwargs,
|
||||
# )
|
||||
|
||||
else:
|
||||
raise Exception(f"Unable to export {config.model_type} model with `--split_modalities`.")
|
||||
|
||||
# Step 2. (optional, recommended) quantize the converted model for fast inference and to reduce model size.
|
||||
if conv_args.quantize:
|
||||
# Update quantize config with model specific defaults
|
||||
quantize_config = MODEL_SPECIFIC_QUANTIZE_PARAMS.get(config.model_type, DEFAULT_QUANTIZE_PARAMS)
|
||||
|
||||
# Update if user specified values
|
||||
if conv_args.per_channel is not None:
|
||||
quantize_config["per_channel"] = conv_args.per_channel
|
||||
|
||||
if conv_args.reduce_range is not None:
|
||||
quantize_config["reduce_range"] = conv_args.reduce_range
|
||||
|
||||
quantize(
|
||||
[
|
||||
os.path.join(output_model_folder, x)
|
||||
for x in os.listdir(output_model_folder)
|
||||
if x.endswith(".onnx") and not x.endswith("_quantized.onnx")
|
||||
],
|
||||
**quantize_config,
|
||||
)
|
||||
|
||||
# Step 3. Move .onnx files to the 'onnx' subfolder
|
||||
os.makedirs(os.path.join(output_model_folder, "onnx"), exist_ok=True)
|
||||
for file in os.listdir(output_model_folder):
|
||||
if file.endswith((".onnx", ".onnx_data")):
|
||||
shutil.move(os.path.join(output_model_folder, file), os.path.join(output_model_folder, "onnx", file))
|
||||
|
||||
# Step 4. Update the generation config if necessary
|
||||
if config.model_type == "whisper":
|
||||
from transformers import GenerationConfig
|
||||
|
||||
from .extra.whisper import get_alignment_heads
|
||||
|
||||
generation_config = GenerationConfig.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
generation_config.alignment_heads = get_alignment_heads(config)
|
||||
generation_config.save_pretrained(output_model_folder)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,29 +0,0 @@
|
||||
import numpy as np
|
||||
import onnx
|
||||
import onnxruntime
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
model_id = "sentence-transformers/paraphrase-MiniLM-L6-v2"
|
||||
output_dir = f"models/{model_id.replace('/', '_')}"
|
||||
model_kwargs = {"output_attentions": True, "return_dict": True}
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
|
||||
model_path = f"{output_dir}/model.onnx"
|
||||
onnx_model = onnx.load(model_path)
|
||||
ort_session = onnxruntime.InferenceSession(model_path)
|
||||
text = "This is a test sentence"
|
||||
tokenizer_output = tokenizer(text, return_tensors="np")
|
||||
input_ids = tokenizer_output["input_ids"]
|
||||
attention_mask = tokenizer_output["attention_mask"]
|
||||
print(attention_mask)
|
||||
# Prepare the input
|
||||
input_ids = np.array(input_ids).astype(np.int64) # Replace your_input_ids with actual input data
|
||||
|
||||
# Run the ONNX model
|
||||
outputs = ort_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
|
||||
|
||||
# Get the attention weights
|
||||
attentions = outputs[-1]
|
||||
|
||||
# Print the attention weights for the first layer and first head
|
||||
print(attentions[0][0])
|
||||
@@ -0,0 +1,56 @@
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import click
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from optimum.onnxruntime import ORTModelForFeatureExtraction
|
||||
from optimum.pipelines import pipeline
|
||||
from torch import Tensor
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
|
||||
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
|
||||
last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
|
||||
return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
|
||||
|
||||
|
||||
def hf_embed(model_id: str, texts: List[str], tokenizer):
|
||||
# Tokenize the input texts
|
||||
model = AutoModel.from_pretrained(model_id)
|
||||
model.eval()
|
||||
encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
|
||||
|
||||
model_output = model(**encoded_input)
|
||||
sentence_embeddings = model_output[0][:, 0]
|
||||
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
|
||||
return sentence_embeddings
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.option("--model_id", help="model id from huggingface.co/models")
|
||||
@click.option("--model_dir", help="The person to greet.")
|
||||
def setup(model_id, model_dir):
|
||||
text = "This is a test sentence"
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
output_dir = Path(model_dir)
|
||||
model = ORTModelForFeatureExtraction.from_pretrained(output_dir)
|
||||
onnx_quant_embed = pipeline(
|
||||
"feature-extraction", model=model, accelerator="ort", tokenizer=tokenizer, return_tensors=True
|
||||
)
|
||||
quant_embeddings = onnx_quant_embed([text])
|
||||
quant_embeddings = F.normalize(quant_embeddings[0][:,0], p=2, dim=1)
|
||||
quant_embeddings = quant_embeddings.detach().numpy()
|
||||
print(quant_embeddings.shape)
|
||||
|
||||
torch_embeddings = hf_embed(model_id, texts=[text], tokenizer=tokenizer)
|
||||
torch_embeddings = F.normalize(torch_embeddings, p=2, dim=1)
|
||||
torch_embeddings = torch_embeddings.detach().numpy()
|
||||
print(torch_embeddings.shape)
|
||||
assert quant_embeddings.shape == torch_embeddings.shape
|
||||
print(np.allclose(quant_embeddings, torch_embeddings, atol=1e-5))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
setup()
|
||||
+2
-19
@@ -1,20 +1,3 @@
|
||||
import importlib.metadata
|
||||
from fastembed.text.text_embedding import TextEmbedding
|
||||
|
||||
from fastembed.image import ImageEmbedding
|
||||
from fastembed.text import TextEmbedding
|
||||
from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
|
||||
from fastembed.late_interaction import LateInteractionTextEmbedding
|
||||
|
||||
try:
|
||||
version = importlib.metadata.version("fastembed")
|
||||
except importlib.metadata.PackageNotFoundError as _:
|
||||
version = importlib.metadata.version("fastembed-gpu")
|
||||
|
||||
__version__ = version
|
||||
__all__ = [
|
||||
"TextEmbedding",
|
||||
"SparseTextEmbedding",
|
||||
"SparseEmbedding",
|
||||
"ImageEmbedding",
|
||||
"LateInteractionTextEmbedding",
|
||||
]
|
||||
__all__ = ["TextEmbedding"]
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
from fastembed.common.types import OnnxProvider, ImageInput, PathInput
|
||||
|
||||
__all__ = ["OnnxProvider", "ImageInput", "PathInput"]
|
||||
|
||||
@@ -11,36 +11,24 @@ from tqdm import tqdm
|
||||
from loguru import logger
|
||||
|
||||
|
||||
def locate_model_file(model_dir: Path, file_names: List[str]) -> Path:
|
||||
"""
|
||||
Find model path for both TransformerJS style `onnx` subdirectory structure and direct model weights structure used
|
||||
by Optimum and Qdrant
|
||||
"""
|
||||
if not model_dir.is_dir():
|
||||
raise ValueError(f"Provided model path '{model_dir}' is not a directory.")
|
||||
|
||||
for file_name in file_names:
|
||||
file_paths = [path for path in model_dir.rglob(file_name) if path.is_file()]
|
||||
|
||||
if file_paths:
|
||||
return file_paths[0]
|
||||
|
||||
raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")
|
||||
|
||||
|
||||
class ModelManagement:
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@classmethod
|
||||
def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Gets the model description from the model_name.
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model.
|
||||
|
||||
raises:
|
||||
ValueError: If the model_name is not supported.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: The model description.
|
||||
"""
|
||||
for model in cls.list_supported_models():
|
||||
if model_name.lower() == model["model"].lower():
|
||||
return model
|
||||
|
||||
raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
|
||||
|
||||
@classmethod
|
||||
def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
|
||||
"""
|
||||
@@ -75,9 +63,7 @@ class ModelManagement:
|
||||
|
||||
show_progress = total_size_in_bytes and show_progress
|
||||
|
||||
with tqdm(
|
||||
total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress
|
||||
) as progress_bar:
|
||||
with tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress) as progress_bar:
|
||||
with open(output_path, "wb") as file:
|
||||
for chunk in response.iter_content(chunk_size=1024):
|
||||
if chunk: # Filter out keep-alive new chunks
|
||||
@@ -86,38 +72,20 @@ class ModelManagement:
|
||||
return output_path
|
||||
|
||||
@classmethod
|
||||
def download_files_from_huggingface(
|
||||
cls,
|
||||
hf_source_repo: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
extra_patterns: Optional[List[str]] = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
def download_files_from_huggingface(cls, hf_source_repo: str, cache_dir: Optional[str] = None) -> str:
|
||||
"""
|
||||
Downloads a model from HuggingFace Hub.
|
||||
Args:
|
||||
hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
|
||||
cache_dir (Optional[str]): The path to the cache directory.
|
||||
extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
|
||||
includes the required model files.
|
||||
Returns:
|
||||
Path: The path to the model directory.
|
||||
"""
|
||||
allow_patterns = [
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"tokenizer_config.json",
|
||||
"special_tokens_map.json",
|
||||
"preprocessor_config.json",
|
||||
]
|
||||
if extra_patterns is not None:
|
||||
allow_patterns.extend(extra_patterns)
|
||||
|
||||
return snapshot_download(
|
||||
repo_id=hf_source_repo,
|
||||
allow_patterns=allow_patterns,
|
||||
ignore_patterns=["model.safetensors", "pytorch_model.bin"],
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=kwargs.get("local_files_only", False),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -174,9 +142,6 @@ class ModelManagement:
|
||||
|
||||
model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
|
||||
|
||||
if model_tar_gz.exists():
|
||||
model_tar_gz.unlink()
|
||||
|
||||
cls.download_file_from_gcs(
|
||||
source_url,
|
||||
output_path=str(model_tar_gz),
|
||||
@@ -192,7 +157,7 @@ class ModelManagement:
|
||||
return model_dir
|
||||
|
||||
@classmethod
|
||||
def download_model(cls, model: Dict[str, Any], cache_dir: Path, **kwargs) -> Path:
|
||||
def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
|
||||
"""
|
||||
Downloads a model from HuggingFace Hub or Google Cloud Storage.
|
||||
|
||||
@@ -221,23 +186,10 @@ class ModelManagement:
|
||||
url_source = model.get("sources", {}).get("url")
|
||||
|
||||
if hf_source:
|
||||
extra_patterns = [model["model_file"]]
|
||||
extra_patterns.extend(model.get("additional_files", []))
|
||||
|
||||
try:
|
||||
return Path(
|
||||
cls.download_files_from_huggingface(
|
||||
hf_source,
|
||||
cache_dir=str(cache_dir),
|
||||
extra_patterns=extra_patterns,
|
||||
local_files_only=kwargs.get("local_files_only", False),
|
||||
)
|
||||
)
|
||||
return Path(cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir)))
|
||||
except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
|
||||
logger.error(
|
||||
f"Could not download model from HuggingFace: {e}"
|
||||
"Falling back to other sources."
|
||||
)
|
||||
logger.error(f"Could not download model from HuggingFace: {e}" "Falling back to other sources.")
|
||||
|
||||
if url_source:
|
||||
return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
|
||||
|
||||
@@ -1,24 +1,11 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
from tokenizers import Tokenizer, AddedToken
|
||||
|
||||
from fastembed.image.transform.operators import Compose
|
||||
|
||||
|
||||
def load_special_tokens(model_dir: Path) -> dict:
|
||||
tokens_map_path = model_dir / "special_tokens_map.json"
|
||||
if not tokens_map_path.exists():
|
||||
raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
|
||||
|
||||
with open(str(tokens_map_path)) as tokens_map_file:
|
||||
tokens_map = json.load(tokens_map_file)
|
||||
|
||||
return tokens_map
|
||||
|
||||
|
||||
def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, dict]:
|
||||
def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
|
||||
config_path = model_dir / "config.json"
|
||||
if not config_path.exists():
|
||||
raise ValueError(f"Could not find config.json in {model_dir}")
|
||||
@@ -31,19 +18,22 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, d
|
||||
if not tokenizer_config_path.exists():
|
||||
raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
|
||||
|
||||
tokens_map_path = model_dir / "special_tokens_map.json"
|
||||
if not tokens_map_path.exists():
|
||||
raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
|
||||
|
||||
with open(str(config_path)) as config_file:
|
||||
config = json.load(config_file)
|
||||
|
||||
with open(str(tokenizer_config_path)) as tokenizer_config_file:
|
||||
tokenizer_config = json.load(tokenizer_config_file)
|
||||
|
||||
tokens_map = load_special_tokens(model_dir)
|
||||
with open(str(tokens_map_path)) as tokens_map_file:
|
||||
tokens_map = json.load(tokens_map_file)
|
||||
|
||||
tokenizer = Tokenizer.from_file(str(tokenizer_path))
|
||||
tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
|
||||
tokenizer.enable_padding(
|
||||
pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
|
||||
)
|
||||
tokenizer.enable_padding(pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"])
|
||||
|
||||
for token in tokens_map.values():
|
||||
if isinstance(token, str):
|
||||
@@ -51,24 +41,12 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, d
|
||||
elif isinstance(token, dict):
|
||||
tokenizer.add_special_tokens([AddedToken(**token)])
|
||||
|
||||
special_token_to_id = {}
|
||||
|
||||
for token in tokens_map.values():
|
||||
if isinstance(token, str):
|
||||
special_token_to_id[token] = tokenizer.token_to_id(token)
|
||||
elif isinstance(token, dict):
|
||||
token_str = token.get("content", "")
|
||||
special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
|
||||
|
||||
return tokenizer, special_token_to_id
|
||||
return tokenizer
|
||||
|
||||
|
||||
def load_preprocessor(model_dir: Path) -> Compose:
|
||||
preprocessor_config_path = model_dir / "preprocessor_config.json"
|
||||
if not preprocessor_config_path.exists():
|
||||
raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")
|
||||
|
||||
with open(str(preprocessor_config_path)) as preprocessor_config_file:
|
||||
preprocessor_config = json.load(preprocessor_config_file)
|
||||
transforms = Compose.from_config(preprocessor_config)
|
||||
return transforms
|
||||
def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
|
||||
# Calculate the Lp norm along the specified dimension
|
||||
norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
|
||||
norm = np.maximum(norm, eps) # Avoid division by zero
|
||||
normalized_array = input_array / norm
|
||||
return normalized_array
|
||||
@@ -1,113 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Generic, Iterable, Optional, Tuple, Type, TypeVar, Sequence
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
from fastembed.common.types import OnnxProvider
|
||||
from fastembed.parallel_processor import Worker
|
||||
|
||||
|
||||
# Holds type of the embedding result
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass
|
||||
class OnnxOutputContext:
|
||||
model_output: np.ndarray
|
||||
attention_mask: Optional[np.ndarray] = None
|
||||
input_ids: Optional[np.ndarray] = None
|
||||
|
||||
|
||||
class OnnxModel(Generic[T]):
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
|
||||
def _preprocess_onnx_input(
|
||||
self, onnx_input: Dict[str, np.ndarray], **kwargs
|
||||
) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Preprocess the onnx input.
|
||||
"""
|
||||
return onnx_input
|
||||
|
||||
def load_onnx_model(
|
||||
self,
|
||||
model_dir: Path,
|
||||
model_file: str,
|
||||
threads: Optional[int],
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
) -> None:
|
||||
model_path = model_dir / model_file
|
||||
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
|
||||
|
||||
onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
|
||||
available_providers = ort.get_available_providers()
|
||||
requested_provider_names = []
|
||||
for provider in onnx_providers:
|
||||
# check providers available
|
||||
provider_name = provider if isinstance(provider, str) else provider[0]
|
||||
requested_provider_names.append(provider_name)
|
||||
if provider_name not in available_providers:
|
||||
raise ValueError(
|
||||
f"Provider {provider_name} is not available. Available providers: {available_providers}"
|
||||
)
|
||||
|
||||
so = ort.SessionOptions()
|
||||
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
if threads is not None:
|
||||
so.intra_op_num_threads = threads
|
||||
so.inter_op_num_threads = threads
|
||||
|
||||
self.model = ort.InferenceSession(
|
||||
str(model_path), providers=onnx_providers, sess_options=so
|
||||
)
|
||||
if "CUDAExecutionProvider" in requested_provider_names:
|
||||
current_providers = self.model.get_providers()
|
||||
if "CUDAExecutionProvider" not in current_providers:
|
||||
warnings.warn(
|
||||
f"Attempt to set CUDAExecutionProvider failed. Current providers: {current_providers}."
|
||||
"If you are using CUDA 12.x, install onnxruntime-gpu via "
|
||||
"`pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/`",
|
||||
RuntimeWarning,
|
||||
)
|
||||
|
||||
def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
|
||||
class EmbeddingWorker(Worker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
) -> OnnxModel:
|
||||
raise NotImplementedError()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
):
|
||||
self.model = self.init_embedding(model_name, cache_dir)
|
||||
|
||||
@classmethod
|
||||
def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
|
||||
return cls(
|
||||
model_name=model_name,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
|
||||
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
@@ -1,14 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
from typing import Union, Iterable, Tuple, Dict, Any
|
||||
|
||||
if sys.version_info >= (3, 10):
|
||||
from typing import TypeAlias
|
||||
else:
|
||||
from typing_extensions import TypeAlias
|
||||
|
||||
|
||||
PathInput: TypeAlias = Union[str, os.PathLike]
|
||||
ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
|
||||
|
||||
OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]
|
||||
@@ -4,16 +4,6 @@ from itertools import islice
|
||||
from pathlib import Path
|
||||
from typing import Union, Iterable, Generator, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
|
||||
# Calculate the Lp norm along the specified dimension
|
||||
norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
|
||||
norm = np.maximum(norm, eps) # Avoid division by zero
|
||||
normalized_array = input_array / norm
|
||||
return normalized_array
|
||||
|
||||
|
||||
def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
|
||||
"""
|
||||
|
||||
@@ -2,12 +2,9 @@ from typing import Optional
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from fastembed import TextEmbedding
|
||||
from fastembed.text.text_embedding import TextEmbedding
|
||||
|
||||
logger.warning(
|
||||
"DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated."
|
||||
"Use from fastembed import TextEmbedding instead."
|
||||
)
|
||||
logger.warning("DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." "Use from fastembed import TextEmbedding instead.")
|
||||
|
||||
DefaultEmbedding = TextEmbedding
|
||||
FlagEmbedding = TextEmbedding
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
from fastembed.image.image_embedding import ImageEmbedding
|
||||
|
||||
|
||||
__all__ = ["ImageEmbedding"]
|
||||
|
||||
@@ -1,87 +0,0 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional, Type, Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common import ImageInput, OnnxProvider
|
||||
from fastembed.image.image_embedding_base import ImageEmbeddingBase
|
||||
from fastembed.image.onnx_embedding import OnnxImageEmbedding
|
||||
|
||||
|
||||
class ImageEmbedding(ImageEmbeddingBase):
|
||||
EMBEDDINGS_REGISTRY: List[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
|
||||
Example:
|
||||
```
|
||||
[
|
||||
{
|
||||
"model": "Qdrant/clip-ViT-B-32-vision",
|
||||
"dim": 512,
|
||||
"description": "CLIP vision encoder based on ViT-B/32",
|
||||
"size_in_GB": 0.33,
|
||||
"sources": {
|
||||
"hf": "Qdrant/clip-ViT-B-32-vision",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
}
|
||||
]
|
||||
```
|
||||
"""
|
||||
result = []
|
||||
for embedding in cls.EMBEDDINGS_REGISTRY:
|
||||
result.extend(embedding.list_supported_models())
|
||||
return result
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
|
||||
supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
|
||||
if any(model_name.lower() == model["model"].lower() for model in supported_models):
|
||||
self.model = EMBEDDING_MODEL_TYPE(
|
||||
model_name, cache_dir, threads, providers=providers, **kwargs
|
||||
)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
f"Model {model_name} is not supported in TextEmbedding."
|
||||
"Please check the supported models using `TextEmbedding.list_supported_models()`"
|
||||
)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
images: ImageInput,
|
||||
batch_size: int = 16,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Encode a list of documents into list of embeddings.
|
||||
We use mean pooling with attention so that the model can handle variable-length inputs.
|
||||
|
||||
Args:
|
||||
images: Iterator of image paths or single image path to embed
|
||||
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
|
||||
parallel:
|
||||
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
||||
If 0, use all available cores.
|
||||
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
||||
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self.model.embed(images, batch_size, parallel, **kwargs)
|
||||
@@ -1,39 +0,0 @@
|
||||
from typing import Iterable, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.model_management import ModelManagement
|
||||
from fastembed.common.types import ImageInput
|
||||
|
||||
|
||||
class ImageEmbeddingBase(ModelManagement):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.cache_dir = cache_dir
|
||||
self.threads = threads
|
||||
self._local_files_only = kwargs.pop("local_files_only", False)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
images: ImageInput,
|
||||
batch_size: int = 16,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Embeds a list of images into a list of embeddings.
|
||||
|
||||
Args:
|
||||
images - The list of image paths to preprocess and embed.
|
||||
**kwargs: Additional keyword argument to pass to the embed method.
|
||||
|
||||
Yields:
|
||||
Iterable[np.ndarray]: The embeddings.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
@@ -1,134 +0,0 @@
|
||||
from typing import Dict, Optional, Iterable, Type, List, Any, Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.onnx_model import OnnxOutputContext
|
||||
from fastembed.common.utils import normalize, define_cache_dir
|
||||
from fastembed.common import ImageInput, OnnxProvider
|
||||
from fastembed.image.image_embedding_base import ImageEmbeddingBase
|
||||
from fastembed.image.onnx_image_model import OnnxImageModel, ImageEmbeddingWorker
|
||||
|
||||
supported_onnx_models = [
|
||||
{
|
||||
"model": "Qdrant/clip-ViT-B-32-vision",
|
||||
"dim": 512,
|
||||
"description": "CLIP vision encoder based on ViT-B/32",
|
||||
"size_in_GB": 0.34,
|
||||
"sources": {
|
||||
"hf": "Qdrant/clip-ViT-B-32-vision",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "Qdrant/resnet50-onnx",
|
||||
"dim": 2048,
|
||||
"description": "ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`__.",
|
||||
"size_in_GB": 0.1,
|
||||
"sources": {
|
||||
"hf": "Qdrant/resnet50-onnx",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
model_name (str): The name of the model to use.
|
||||
cache_dir (str, optional): The path to the cache directory.
|
||||
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
||||
Defaults to `fastembed_cache` in the system's temp directory.
|
||||
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
|
||||
"""
|
||||
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
model_description = self._get_model_description(model_name)
|
||||
cache_dir = define_cache_dir(cache_dir)
|
||||
model_dir = self.download_model(
|
||||
model_description, cache_dir, local_files_only=self._local_files_only
|
||||
)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
providers=providers,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
"""
|
||||
return supported_onnx_models
|
||||
|
||||
def embed(
|
||||
self,
|
||||
images: ImageInput,
|
||||
batch_size: int = 16,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Encode a list of images into list of embeddings.
|
||||
We use mean pooling with attention so that the model can handle variable-length inputs.
|
||||
|
||||
Args:
|
||||
images: Iterator of image paths or single image path to embed
|
||||
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
|
||||
parallel:
|
||||
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
||||
If 0, use all available cores.
|
||||
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
||||
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self._embed_images(
|
||||
model_name=self.model_name,
|
||||
cache_dir=str(self.cache_dir),
|
||||
images=images,
|
||||
batch_size=batch_size,
|
||||
parallel=parallel,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
|
||||
return OnnxImageEmbeddingWorker
|
||||
|
||||
def _preprocess_onnx_input(
|
||||
self, onnx_input: Dict[str, np.ndarray], **kwargs
|
||||
) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Preprocess the onnx input.
|
||||
"""
|
||||
|
||||
return onnx_input
|
||||
|
||||
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
|
||||
return normalize(output.model_output).astype(np.float32)
|
||||
|
||||
|
||||
class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
) -> OnnxImageEmbedding:
|
||||
return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
|
||||
@@ -1,107 +0,0 @@
|
||||
import os
|
||||
import contextlib
|
||||
from multiprocessing import get_all_start_methods
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Sequence
|
||||
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.preprocessor_utils import load_preprocessor
|
||||
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
|
||||
from fastembed.common import PathInput, ImageInput, OnnxProvider
|
||||
from fastembed.common.utils import iter_batch
|
||||
from fastembed.parallel_processor import ParallelWorkerPool
|
||||
|
||||
# Holds type of the embedding result
|
||||
|
||||
|
||||
class OnnxImageModel(OnnxModel[T]):
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.processor = None
|
||||
|
||||
def _preprocess_onnx_input(
|
||||
self, onnx_input: Dict[str, np.ndarray], **kwargs
|
||||
) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Preprocess the onnx input.
|
||||
"""
|
||||
return onnx_input
|
||||
|
||||
def load_onnx_model(
|
||||
self,
|
||||
model_dir: Path,
|
||||
model_file: str,
|
||||
threads: Optional[int],
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
) -> None:
|
||||
super().load_onnx_model(
|
||||
model_dir=model_dir, model_file=model_file, threads=threads, providers=providers
|
||||
)
|
||||
self.processor = load_preprocessor(model_dir=model_dir)
|
||||
|
||||
def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
|
||||
return {node.name: encoded for node in self.model.get_inputs()}
|
||||
|
||||
def onnx_embed(self, images: List[PathInput], **kwargs) -> OnnxOutputContext:
|
||||
with contextlib.ExitStack():
|
||||
image_files = [Image.open(image) for image in images]
|
||||
encoded = self.processor(image_files)
|
||||
onnx_input = self._build_onnx_input(encoded)
|
||||
onnx_input = self._preprocess_onnx_input(onnx_input)
|
||||
model_output = self.model.run(None, onnx_input)
|
||||
embeddings = model_output[0].reshape(len(images), -1)
|
||||
return OnnxOutputContext(
|
||||
model_output=embeddings
|
||||
)
|
||||
|
||||
def _embed_images(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
images: ImageInput,
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
) -> Iterable[T]:
|
||||
is_small = False
|
||||
|
||||
if isinstance(images, str) or isinstance(images, Path):
|
||||
images = [images]
|
||||
is_small = True
|
||||
|
||||
if isinstance(images, list):
|
||||
if len(images) < batch_size:
|
||||
is_small = True
|
||||
|
||||
if parallel == 0:
|
||||
parallel = os.cpu_count()
|
||||
|
||||
if parallel is None or is_small:
|
||||
for batch in iter_batch(images, batch_size):
|
||||
yield from self._post_process_onnx_output(self.onnx_embed(batch))
|
||||
else:
|
||||
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
|
||||
params = {
|
||||
"model_name": model_name,
|
||||
"cache_dir": cache_dir,
|
||||
}
|
||||
pool = ParallelWorkerPool(
|
||||
parallel, self._get_worker_class(), start_method=start_method
|
||||
)
|
||||
for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
|
||||
yield from self._post_process_onnx_output(batch)
|
||||
|
||||
|
||||
class ImageEmbeddingWorker(EmbeddingWorker):
|
||||
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
|
||||
for idx, batch in items:
|
||||
embeddings = self.model.onnx_embed(batch)
|
||||
yield idx, embeddings
|
||||
@@ -1,125 +0,0 @@
|
||||
from typing import Union, Tuple, Sized
|
||||
|
||||
from PIL import Image
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def convert_to_rgb(image: Image.Image) -> Image.Image:
|
||||
if image.mode == "RGB":
|
||||
return image
|
||||
|
||||
image = image.convert("RGB")
|
||||
return image
|
||||
|
||||
|
||||
def center_crop(
|
||||
image: Union[Image.Image, np.ndarray],
|
||||
size: Tuple[int, int],
|
||||
) -> np.ndarray:
|
||||
if isinstance(image, np.ndarray):
|
||||
_, orig_height, orig_width = image.shape
|
||||
else:
|
||||
orig_height, orig_width = image.height, image.width
|
||||
# (H, W, C) -> (C, H, W)
|
||||
image = np.array(image).transpose((2, 0, 1))
|
||||
|
||||
crop_height, crop_width = size
|
||||
|
||||
# left upper corner (0, 0)
|
||||
top = (orig_height - crop_height) // 2
|
||||
bottom = top + crop_height
|
||||
left = (orig_width - crop_width) // 2
|
||||
right = left + crop_width
|
||||
|
||||
# Check if cropped area is within image boundaries
|
||||
if top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width:
|
||||
image = image[..., top:bottom, left:right]
|
||||
return image
|
||||
|
||||
# Padding with zeros
|
||||
new_height = max(crop_height, orig_height)
|
||||
new_width = max(crop_width, orig_width)
|
||||
new_shape = image.shape[:-2] + (new_height, new_width)
|
||||
new_image = np.zeros_like(image, shape=new_shape)
|
||||
|
||||
top_pad = (new_height - orig_height) // 2
|
||||
bottom_pad = top_pad + orig_height
|
||||
left_pad = (new_width - orig_width) // 2
|
||||
right_pad = left_pad + orig_width
|
||||
new_image[..., top_pad:bottom_pad, left_pad:right_pad] = image
|
||||
|
||||
top += top_pad
|
||||
bottom += top_pad
|
||||
left += left_pad
|
||||
right += left_pad
|
||||
|
||||
new_image = new_image[
|
||||
..., max(0, top) : min(new_height, bottom), max(0, left) : min(new_width, right)
|
||||
]
|
||||
|
||||
return new_image
|
||||
|
||||
|
||||
def normalize(
|
||||
image: np.ndarray,
|
||||
mean=Union[float, np.ndarray],
|
||||
std=Union[float, np.ndarray],
|
||||
) -> np.ndarray:
|
||||
if not isinstance(image, np.ndarray):
|
||||
raise ValueError("image must be a numpy array")
|
||||
|
||||
num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0]
|
||||
|
||||
if not np.issubdtype(image.dtype, np.floating):
|
||||
image = image.astype(np.float32)
|
||||
|
||||
if isinstance(mean, Sized):
|
||||
if len(mean) != num_channels:
|
||||
raise ValueError(
|
||||
f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}"
|
||||
)
|
||||
else:
|
||||
mean = [mean] * num_channels
|
||||
mean = np.array(mean, dtype=image.dtype)
|
||||
|
||||
if isinstance(std, Sized):
|
||||
if len(std) != num_channels:
|
||||
raise ValueError(
|
||||
f"std must have {num_channels} elements if it is an iterable, got {len(std)}"
|
||||
)
|
||||
else:
|
||||
std = [std] * num_channels
|
||||
std = np.array(std, dtype=image.dtype)
|
||||
|
||||
image = ((image.T - mean) / std).T
|
||||
return image
|
||||
|
||||
|
||||
def resize(
|
||||
image: Image,
|
||||
size: Union[int, Tuple[int, int]],
|
||||
resample: Image.Resampling = Image.Resampling.BILINEAR,
|
||||
) -> Image:
|
||||
if isinstance(size, tuple):
|
||||
return image.resize(size, resample)
|
||||
|
||||
height, width = image.height, image.width
|
||||
short, long = (width, height) if width <= height else (height, width)
|
||||
|
||||
new_short, new_long = size, int(size * long / short)
|
||||
if width <= height:
|
||||
new_size = (new_short, new_long)
|
||||
else:
|
||||
new_size = (new_long, new_short)
|
||||
return image.resize(new_size, resample)
|
||||
|
||||
|
||||
def rescale(image: np.ndarray, scale: float, dtype=np.float32) -> np.ndarray:
|
||||
return (image * scale).astype(dtype)
|
||||
|
||||
|
||||
def pil2ndarray(image: Union[Image.Image, np.ndarray]):
|
||||
if isinstance(image, Image.Image):
|
||||
return np.asarray(image).transpose((2, 0, 1))
|
||||
return image
|
||||
@@ -1,168 +0,0 @@
|
||||
from typing import List, Tuple, Union, Any, Dict
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from fastembed.image.transform.functional import (
|
||||
center_crop,
|
||||
normalize,
|
||||
resize,
|
||||
convert_to_rgb,
|
||||
rescale,
|
||||
pil2ndarray
|
||||
)
|
||||
|
||||
|
||||
class Transform:
|
||||
def __call__(self, images: List) -> Union[List[Image.Image], List[np.ndarray]]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
|
||||
class ConvertToRGB(Transform):
|
||||
def __call__(self, images: List[Image.Image]) -> List[Image.Image]:
|
||||
return [convert_to_rgb(image=image) for image in images]
|
||||
|
||||
|
||||
class CenterCrop(Transform):
|
||||
def __init__(self, size: Tuple[int, int]):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, images: List[Image.Image]) -> List[np.ndarray]:
|
||||
return [center_crop(image=image, size=self.size) for image in images]
|
||||
|
||||
|
||||
class Normalize(Transform):
|
||||
def __init__(self, mean: Union[float, List[float]], std: Union[float, List[float]]):
|
||||
self.mean = mean
|
||||
self.std = std
|
||||
|
||||
def __call__(self, images: List[np.ndarray]) -> List[np.ndarray]:
|
||||
return [normalize(image, mean=self.mean, std=self.std) for image in images]
|
||||
|
||||
|
||||
class Resize(Transform):
|
||||
def __init__(
|
||||
self,
|
||||
size: Union[int, Tuple[int, int]],
|
||||
resample: Image.Resampling = Image.Resampling.BICUBIC,
|
||||
):
|
||||
self.size = size
|
||||
self.resample = resample
|
||||
|
||||
def __call__(self, images: List[Image.Image]) -> List[Image.Image]:
|
||||
return [resize(image, size=self.size, resample=self.resample) for image in images]
|
||||
|
||||
|
||||
class Rescale(Transform):
|
||||
def __init__(self, scale: float = 1 / 255):
|
||||
self.scale = scale
|
||||
|
||||
def __call__(self, images: List[np.ndarray]) -> List[np.ndarray]:
|
||||
return [rescale(image, scale=self.scale) for image in images]
|
||||
|
||||
class PILtoNDarray(Transform):
|
||||
def __call__(self, images: List[Union[Image.Image, np.ndarray]]) -> List[np.ndarray]:
|
||||
return [pil2ndarray(image) for image in images]
|
||||
|
||||
class Compose:
|
||||
def __init__(self, transforms: List[Transform]):
|
||||
self.transforms = transforms
|
||||
|
||||
def __call__(self, images: Union[List[Image.Image], List[np.ndarray]]) -> Union[List[np.ndarray], List[Image.Image]]:
|
||||
for transform in self.transforms:
|
||||
images = transform(images)
|
||||
return images
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> "Compose":
|
||||
"""Creates processor from a config dict.
|
||||
Args:
|
||||
config (Dict[str, Any]): Configuration dictionary.
|
||||
|
||||
Valid keys:
|
||||
- do_resize
|
||||
- size
|
||||
- do_center_crop
|
||||
- crop_size
|
||||
- do_rescale
|
||||
- rescale_factor
|
||||
- do_normalize
|
||||
- image_mean
|
||||
- image_std
|
||||
Valid size keys (nested):
|
||||
- {"height", "width"}
|
||||
- {"shortest_edge"}
|
||||
|
||||
Returns:
|
||||
Compose: Image processor.
|
||||
"""
|
||||
transforms = []
|
||||
cls._get_convert_to_rgb(transforms, config)
|
||||
cls._get_resize(transforms, config)
|
||||
cls._get_center_crop(transforms, config)
|
||||
cls._get_pil2ndarray(transforms, config)
|
||||
cls._get_rescale(transforms, config)
|
||||
cls._get_normalize(transforms, config)
|
||||
return cls(transforms=transforms)
|
||||
|
||||
@staticmethod
|
||||
def _get_convert_to_rgb(transforms: List[Transform], config: Dict[str, Any]):
|
||||
transforms.append(ConvertToRGB())
|
||||
|
||||
@staticmethod
|
||||
def _get_resize(transforms: List[Transform], config: Dict[str, Any]):
|
||||
mode = config.get('image_processor_type', 'CLIPImageProcessor')
|
||||
if mode == 'CLIPImageProcessor':
|
||||
if config.get("do_resize", False):
|
||||
size = config["size"]
|
||||
if "shortest_edge" in size:
|
||||
size = size["shortest_edge"]
|
||||
elif "height" in size and "width" in size:
|
||||
size = (size["height"], size["width"])
|
||||
else:
|
||||
raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")
|
||||
transforms.append(Resize(size=size, resample=config.get("resample", Image.Resampling.BICUBIC)))
|
||||
elif mode == 'ConvNextFeatureExtractor':
|
||||
if 'size' in config and "shortest_edge" not in config['size']:
|
||||
raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}")
|
||||
shortest_edge = config['size']["shortest_edge"]
|
||||
crop_pct = config.get("crop_pct", 0.875)
|
||||
if shortest_edge < 384:
|
||||
# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
|
||||
resize_shortest_edge = int(shortest_edge / crop_pct)
|
||||
transforms.append(Resize(size=resize_shortest_edge, resample=config.get("resample", Image.Resampling.BICUBIC)))
|
||||
transforms.append(CenterCrop(size=(shortest_edge, shortest_edge)))
|
||||
else:
|
||||
transforms.append(Resize(size=(shortest_edge, shortest_edge), resample=config.get("resample", Image.Resampling.BICUBIC)))
|
||||
|
||||
@staticmethod
|
||||
def _get_center_crop(transforms: List[Transform], config: Dict[str, Any]):
|
||||
mode = config.get('image_processor_type', 'CLIPImageProcessor')
|
||||
if mode == 'CLIPImageProcessor':
|
||||
if config.get("do_center_crop", False):
|
||||
crop_size = config["crop_size"]
|
||||
if isinstance(crop_size, int):
|
||||
crop_size = (crop_size, crop_size)
|
||||
elif isinstance(crop_size, dict):
|
||||
crop_size = (crop_size["height"], crop_size["width"])
|
||||
else:
|
||||
raise ValueError(f"Invalid crop size: {crop_size}")
|
||||
transforms.append(CenterCrop(size=crop_size))
|
||||
elif mode == 'ConvNextFeatureExtractor':
|
||||
pass
|
||||
else:
|
||||
raise ValueError(f"Preprocessor {mode} is not supported")
|
||||
|
||||
@staticmethod
|
||||
def _get_pil2ndarray(transforms: List[Transform], config: Dict[str, Any]):
|
||||
transforms.append(PILtoNDarray())
|
||||
|
||||
@staticmethod
|
||||
def _get_rescale(transforms: List[Transform], config: Dict[str, Any]):
|
||||
if config.get("do_rescale", True):
|
||||
rescale_factor = config.get("rescale_factor", 1 / 255)
|
||||
transforms.append(Rescale(scale=rescale_factor))
|
||||
|
||||
@staticmethod
|
||||
def _get_normalize(transforms: List[Transform], config: Dict[str, Any]):
|
||||
if config.get("do_normalize", False):
|
||||
transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))
|
||||
@@ -1,4 +0,0 @@
|
||||
from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding
|
||||
|
||||
|
||||
__all__ = ["LateInteractionTextEmbedding"]
|
||||
@@ -1,196 +0,0 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional, Union, Type, Sequence
|
||||
import string
|
||||
|
||||
import numpy as np
|
||||
from tokenizers import Encoding
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.common.onnx_model import OnnxOutputContext
|
||||
from fastembed.common.utils import define_cache_dir
|
||||
from fastembed.late_interaction.late_interaction_embedding_base import (
|
||||
LateInteractionTextEmbeddingBase,
|
||||
)
|
||||
from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
|
||||
|
||||
|
||||
supported_colbert_models = [
|
||||
{
|
||||
"model": "colbert-ir/colbertv2.0",
|
||||
"dim": 128,
|
||||
"description": "Late interaction model",
|
||||
"size_in_GB": 0.44,
|
||||
"sources": {
|
||||
"hf": "colbert-ir/colbertv2.0",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
|
||||
QUERY_MARKER_TOKEN_ID = 1
|
||||
DOCUMENT_MARKER_TOKEN_ID = 2
|
||||
MIN_QUERY_LENGTH = 32
|
||||
MASK_TOKEN = "[MASK]"
|
||||
|
||||
def _post_process_onnx_output(
|
||||
self, output: OnnxOutputContext, is_doc: bool = True
|
||||
) -> Iterable[np.ndarray]:
|
||||
if not is_doc:
|
||||
return output.model_output.astype(np.float32)
|
||||
|
||||
for i, token_sequence in enumerate(output.input_ids):
|
||||
for j, token_id in enumerate(token_sequence):
|
||||
if token_id in self.skip_list or token_id == self.pad_token_id:
|
||||
output.attention_mask[i, j] = 0
|
||||
|
||||
output.model_output *= np.expand_dims(output.attention_mask, 2).astype(np.float32)
|
||||
norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
|
||||
norm_clamped = np.maximum(norm, 1e-12)
|
||||
output.model_output /= norm_clamped
|
||||
return output.model_output.astype(np.float32)
|
||||
|
||||
def _preprocess_onnx_input(
|
||||
self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True
|
||||
) -> Dict[str, np.ndarray]:
|
||||
if is_doc:
|
||||
onnx_input["input_ids"][:, 1] = self.DOCUMENT_MARKER_TOKEN_ID
|
||||
else:
|
||||
onnx_input["input_ids"][:, 1] = self.QUERY_MARKER_TOKEN_ID
|
||||
return onnx_input
|
||||
|
||||
def tokenize(self, documents: List[str], is_doc: bool = True) -> List[Encoding]:
|
||||
return (
|
||||
self._tokenize_documents(documents=documents)
|
||||
if is_doc
|
||||
else self._tokenize_query(query=next(iter(documents)))
|
||||
)
|
||||
|
||||
def _tokenize_query(self, query: str) -> List[Encoding]:
|
||||
# ". " is added to a query to be replaced with a special query token
|
||||
query = [f". {query}"]
|
||||
encoded = self.tokenizer.encode_batch(query)
|
||||
# colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
|
||||
if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
|
||||
prev_padding = None
|
||||
if self.tokenizer.padding:
|
||||
prev_padding = self.tokenizer.padding
|
||||
self.tokenizer.enable_padding(
|
||||
pad_token=self.MASK_TOKEN, pad_id=self.mask_token_id, length=self.MIN_QUERY_LENGTH
|
||||
)
|
||||
encoded = self.tokenizer.encode_batch(query)
|
||||
if prev_padding is None:
|
||||
self.tokenizer.no_padding()
|
||||
else:
|
||||
self.tokenizer.enable_padding(**prev_padding)
|
||||
return encoded
|
||||
|
||||
def _tokenize_documents(self, documents: List[str]) -> List[Encoding]:
|
||||
# ". " is added to a document to be replaced with a special document token
|
||||
documents = [". " + doc for doc in documents]
|
||||
encoded = self.tokenizer.encode_batch(documents)
|
||||
return encoded
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
"""
|
||||
return supported_colbert_models
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
model_name (str): The name of the model to use.
|
||||
cache_dir (str, optional): The path to the cache directory.
|
||||
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
||||
Defaults to `fastembed_cache` in the system's temp directory.
|
||||
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
|
||||
"""
|
||||
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
model_description = self._get_model_description(model_name)
|
||||
cache_dir = define_cache_dir(cache_dir)
|
||||
|
||||
model_dir = self.download_model(
|
||||
model_description, cache_dir, local_files_only=self._local_files_only
|
||||
)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
providers=providers,
|
||||
)
|
||||
self.mask_token_id = self.special_token_to_id["[MASK]"]
|
||||
self.pad_token_id = self.tokenizer.padding["pad_id"]
|
||||
|
||||
self.skip_list = {
|
||||
self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]
|
||||
for symbol in string.punctuation
|
||||
}
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Encode a list of documents into list of embeddings.
|
||||
We use mean pooling with attention so that the model can handle variable-length inputs.
|
||||
|
||||
Args:
|
||||
documents: Iterator of documents or single document to embed
|
||||
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
|
||||
parallel:
|
||||
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
||||
If 0, use all available cores.
|
||||
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
||||
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self._embed_documents(
|
||||
model_name=self.model_name,
|
||||
cache_dir=str(self.cache_dir),
|
||||
documents=documents,
|
||||
batch_size=batch_size,
|
||||
parallel=parallel,
|
||||
)
|
||||
|
||||
def query_embed(self, query: Union[str, List[str]], **kwargs) -> np.ndarray:
|
||||
if isinstance(query, str):
|
||||
query = [query]
|
||||
|
||||
for text in query:
|
||||
yield from self._post_process_onnx_output(
|
||||
self.onnx_embed([text], is_doc=False), is_doc=False
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
|
||||
return ColbertEmbeddingWorker
|
||||
|
||||
|
||||
class ColbertEmbeddingWorker(TextEmbeddingWorker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
) -> Colbert:
|
||||
return Colbert(model_name=model_name, cache_dir=cache_dir, threads=1)
|
||||
@@ -1,60 +0,0 @@
|
||||
from typing import Iterable, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.model_management import ModelManagement
|
||||
|
||||
|
||||
class LateInteractionTextEmbeddingBase(ModelManagement):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.cache_dir = cache_dir
|
||||
self.threads = threads
|
||||
self._local_files_only = kwargs.pop("local_files_only", False)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
raise NotImplementedError()
|
||||
|
||||
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Embeds a list of text passages into a list of embeddings.
|
||||
|
||||
Args:
|
||||
texts (Iterable[str]): The list of texts to embed.
|
||||
**kwargs: Additional keyword argument to pass to the embed method.
|
||||
|
||||
Yields:
|
||||
Iterable[np.ndarray]: The embeddings.
|
||||
"""
|
||||
|
||||
# This is model-specific, so that different models can have specialized implementations
|
||||
yield from self.embed(texts, **kwargs)
|
||||
|
||||
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Embeds queries
|
||||
|
||||
Args:
|
||||
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
|
||||
|
||||
Returns:
|
||||
Iterable[np.ndarray]: The embeddings.
|
||||
"""
|
||||
|
||||
# This is model-specific, so that different models can have specialized implementations
|
||||
if isinstance(query, str):
|
||||
yield from self.embed([query], **kwargs)
|
||||
if isinstance(query, Iterable):
|
||||
yield from self.embed(query, **kwargs)
|
||||
@@ -1,104 +0,0 @@
|
||||
from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.late_interaction.late_interaction_embedding_base import (
|
||||
LateInteractionTextEmbeddingBase,
|
||||
)
|
||||
from fastembed.late_interaction.colbert import Colbert
|
||||
|
||||
|
||||
class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
|
||||
EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [
|
||||
Colbert,
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
|
||||
Example:
|
||||
```
|
||||
[
|
||||
{
|
||||
"model": "prithvida/SPLADE_PP_en_v1",
|
||||
"vocab_size": 30522,
|
||||
"description": "Independent Implementation of SPLADE++ Model for English",
|
||||
"size_in_GB": 0.532,
|
||||
"sources": {
|
||||
"hf": "qdrant/SPLADE_PP_en_v1",
|
||||
},
|
||||
}
|
||||
]
|
||||
```
|
||||
"""
|
||||
result = []
|
||||
for embedding in cls.EMBEDDINGS_REGISTRY:
|
||||
result.extend(embedding.list_supported_models())
|
||||
return result
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
|
||||
supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
|
||||
if any(model_name.lower() == model["model"].lower() for model in supported_models):
|
||||
self.model = EMBEDDING_MODEL_TYPE(
|
||||
model_name, cache_dir, threads, providers=providers, **kwargs
|
||||
)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
f"Model {model_name} is not supported in SparseTextEmbedding."
|
||||
"Please check the supported models using `SparseTextEmbedding.list_supported_models()`"
|
||||
)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Encode a list of documents into list of embeddings.
|
||||
We use mean pooling with attention so that the model can handle variable-length inputs.
|
||||
|
||||
Args:
|
||||
documents: Iterator of documents or single document to embed
|
||||
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
|
||||
parallel:
|
||||
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
||||
If 0, use all available cores.
|
||||
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
||||
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
|
||||
|
||||
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
Embeds queries
|
||||
|
||||
Args:
|
||||
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
|
||||
|
||||
Returns:
|
||||
Iterable[np.ndarray]: The embeddings.
|
||||
"""
|
||||
|
||||
# This is model-specific, so that different models can have specialized implementations
|
||||
yield from self.model.query_embed(query, **kwargs)
|
||||
@@ -128,9 +128,7 @@ class ParallelWorkerPool:
|
||||
yield buffer.pop(next_expected)
|
||||
next_expected += 1
|
||||
|
||||
def semi_ordered_map(
|
||||
self, stream: Iterable[Any], *args: Any, **kwargs: Any
|
||||
) -> Iterable[Tuple[int, Any]]:
|
||||
def semi_ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Tuple[int, Any]]:
|
||||
try:
|
||||
self.start(**kwargs)
|
||||
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
from fastembed.sparse.sparse_embedding_base import SparseEmbedding
|
||||
from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
|
||||
|
||||
__all__ = ["SparseEmbedding", "SparseTextEmbedding"]
|
||||
|
||||
@@ -1,281 +0,0 @@
|
||||
import math
|
||||
import string
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
|
||||
|
||||
import numpy as np
|
||||
import mmh3
|
||||
from snowballstemmer import stemmer as get_stemmer
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.common.onnx_model import OnnxOutputContext
|
||||
from fastembed.common.utils import define_cache_dir
|
||||
from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
|
||||
from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
|
||||
|
||||
supported_bm42_models = [
|
||||
{
|
||||
"model": "Qdrant/bm42-all-minilm-l6-v2-attentions",
|
||||
"vocab_size": 30522,
|
||||
"description": "Light sparse embedding model, which assigns an importance score to each token in the text",
|
||||
"size_in_GB": 0.09,
|
||||
"sources": {
|
||||
"hf": "Qdrant/all_miniLM_L6_v2_with_attentions",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
"additional_files": ["stopwords.txt"],
|
||||
},
|
||||
]
|
||||
|
||||
MODEL_TO_LANGUAGE = {
|
||||
"Qdrant/bm42-all-minilm-l6-v2-attentions": "english",
|
||||
}
|
||||
|
||||
|
||||
class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
|
||||
"""
|
||||
Bm42 is an extension of BM25, which tries to better evaluate importance of tokens in the documents,
|
||||
by extracting attention weights from the transformer model.
|
||||
|
||||
Traditional BM25 uses a count of tokens in the document to evaluate the importance of the token,
|
||||
but this approach doesn't work well with short documents or chunks of text, as almost all tokens
|
||||
there are unique.
|
||||
|
||||
BM42 addresses this issue by replacing the token count with the attention weights from the transformer model.
|
||||
This allows sparse embeddings to work well with short documents, handle rare tokens and leverage traditional NLP
|
||||
techniques like stemming and stopwords.
|
||||
|
||||
WARNING: This model is expected to be used with `modifier="idf"` in the sparse vector index of Qdrant.
|
||||
"""
|
||||
|
||||
ONNX_OUTPUT_NAMES = ["attention_6"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
alpha: float = 0.5,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
model_name (str): The name of the model to use.
|
||||
cache_dir (str, optional): The path to the cache directory.
|
||||
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
||||
Defaults to `fastembed_cache` in the system's temp directory.
|
||||
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
|
||||
providers (Optional[Sequence[OnnxProvider]], optional): The providers to use for onnxruntime.
|
||||
alpha (float, optional): Parameter, that defines the importance of the token weight in the document
|
||||
versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
|
||||
It is recommended to only change this parameter based on training data for a specific dataset.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
|
||||
"""
|
||||
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
model_description = self._get_model_description(model_name)
|
||||
cache_dir = define_cache_dir(cache_dir)
|
||||
|
||||
model_dir = self.download_model(
|
||||
model_description, cache_dir, local_files_only=self._local_files_only
|
||||
)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
providers=providers,
|
||||
)
|
||||
|
||||
self.invert_vocab = {}
|
||||
|
||||
for token, idx in self.tokenizer.get_vocab().items():
|
||||
self.invert_vocab[idx] = token
|
||||
|
||||
self.special_tokens = set(self.special_token_to_id.keys())
|
||||
self.special_tokens_ids = set(self.special_token_to_id.values())
|
||||
self.punctuation = set(string.punctuation)
|
||||
self.stopwords = set(self._load_stopwords(model_dir))
|
||||
self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
|
||||
self.alpha = alpha
|
||||
|
||||
def _filter_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
|
||||
result = []
|
||||
for token, value in tokens:
|
||||
if token in self.stopwords or token in self.punctuation:
|
||||
continue
|
||||
result.append((token, value))
|
||||
return result
|
||||
|
||||
def _stem_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
|
||||
result = []
|
||||
for token, value in tokens:
|
||||
processed_token = self.stemmer.stemWord(token)
|
||||
result.append((processed_token, value))
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def _aggregate_weights(cls, tokens: List[Tuple[str, List[int]]], weights: List[float]) -> List[Tuple[str, float]]:
|
||||
result = []
|
||||
for token, idxs in tokens:
|
||||
sum_weight = sum(weights[idx] for idx in idxs)
|
||||
result.append((token, sum_weight))
|
||||
return result
|
||||
|
||||
def _reconstruct_bpe(
|
||||
self, bpe_tokens: Iterable[Tuple[int, str]]
|
||||
) -> List[Tuple[str, List[int]]]:
|
||||
result = []
|
||||
acc = ""
|
||||
acc_idx = []
|
||||
|
||||
continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix
|
||||
continuing_subword_prefix_len = len(continuing_subword_prefix)
|
||||
|
||||
for idx, token in bpe_tokens:
|
||||
if token in self.special_tokens:
|
||||
continue
|
||||
|
||||
if token.startswith(continuing_subword_prefix):
|
||||
acc += token[continuing_subword_prefix_len:]
|
||||
acc_idx.append(idx)
|
||||
else:
|
||||
if acc:
|
||||
result.append((acc, acc_idx))
|
||||
acc_idx = []
|
||||
acc = token
|
||||
acc_idx.append(idx)
|
||||
|
||||
if acc:
|
||||
result.append((acc, acc_idx))
|
||||
|
||||
return result
|
||||
|
||||
def _rescore_vector(self, vector: Dict[str, float]) -> Dict[int, float]:
|
||||
"""
|
||||
Orders all tokens in the vector by their importance and generates a new score based on the importance order.
|
||||
So that the scoring doesn't depend on absolute values assigned by the model, but on the relative importance.
|
||||
"""
|
||||
|
||||
new_vector = {}
|
||||
|
||||
for token, value in vector.items():
|
||||
token_id = abs(mmh3.hash(token))
|
||||
# Examples:
|
||||
# Num 0: Log(1/1 + 1) = 0.6931471805599453
|
||||
# Num 1: Log(1/2 + 1) = 0.4054651081081644
|
||||
# Num 2: Log(1/3 + 1) = 0.28768207245178085
|
||||
new_vector[token_id] = math.log(1. + value) ** self.alpha # value
|
||||
|
||||
return new_vector
|
||||
|
||||
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
|
||||
token_ids_batch = output.input_ids
|
||||
|
||||
# attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
|
||||
pooled_attention = np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask
|
||||
|
||||
for document_token_ids, attention_value in zip(token_ids_batch, pooled_attention):
|
||||
document_tokens_with_ids = (
|
||||
(idx, self.invert_vocab[token_id])
|
||||
for idx, token_id in enumerate(document_token_ids)
|
||||
)
|
||||
|
||||
reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
|
||||
|
||||
filtered = self._filter_pair_tokens(reconstructed)
|
||||
|
||||
stemmed = self._stem_pair_tokens(filtered)
|
||||
|
||||
weighted = self._aggregate_weights(stemmed, attention_value)
|
||||
|
||||
max_token_weight = {}
|
||||
|
||||
for token, weight in weighted:
|
||||
max_token_weight[token] = max(max_token_weight.get(token, 0), weight)
|
||||
|
||||
rescored = self._rescore_vector(max_token_weight)
|
||||
|
||||
yield SparseEmbedding.from_dict(rescored)
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
"""
|
||||
return supported_bm42_models
|
||||
|
||||
@classmethod
|
||||
def _load_stopwords(cls, model_dir: Path) -> List[str]:
|
||||
stopwords_path = model_dir / "stopwords.txt"
|
||||
if not stopwords_path.exists():
|
||||
return []
|
||||
|
||||
with open(stopwords_path, "r") as f:
|
||||
return f.read().splitlines()
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[SparseEmbedding]:
|
||||
"""
|
||||
Encode a list of documents into list of embeddings.
|
||||
We use mean pooling with attention so that the model can handle variable-length inputs.
|
||||
|
||||
Args:
|
||||
documents: Iterator of documents or single document to embed
|
||||
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
|
||||
parallel:
|
||||
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
||||
If 0, use all available cores.
|
||||
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
||||
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self._embed_documents(
|
||||
model_name=self.model_name,
|
||||
cache_dir=str(self.cache_dir),
|
||||
documents=documents,
|
||||
batch_size=batch_size,
|
||||
parallel=parallel,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _query_rehash(cls, tokens: Iterable[str]) -> Dict[int, float]:
|
||||
result = {}
|
||||
for token in tokens:
|
||||
token_id = abs(mmh3.hash(token))
|
||||
result[token_id] = 1.0
|
||||
return result
|
||||
|
||||
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
|
||||
"""
|
||||
To emulate BM25 behaviour, we don't need to use smart weights in the query, and
|
||||
it's enough to just hash the tokens and assign a weight of 1.0 to them.
|
||||
It is also faster, as we don't need to run the model for the query.
|
||||
"""
|
||||
if isinstance(query, str):
|
||||
query = [query]
|
||||
|
||||
for text in query:
|
||||
encoded = self.tokenizer.encode(text)
|
||||
document_tokens_with_ids = enumerate(encoded.tokens)
|
||||
reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
|
||||
filtered = self._filter_pair_tokens(reconstructed)
|
||||
stemmed = self._stem_pair_tokens(filtered)
|
||||
|
||||
yield SparseEmbedding.from_dict(self._query_rehash(token for token, _ in stemmed))
|
||||
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
|
||||
return TextEmbeddingWorker
|
||||
@@ -1,81 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Iterable, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.model_management import ModelManagement
|
||||
|
||||
|
||||
@dataclass
|
||||
class SparseEmbedding:
|
||||
values: np.ndarray
|
||||
indices: np.ndarray
|
||||
|
||||
def as_object(self) -> Dict[str, np.ndarray]:
|
||||
return {
|
||||
"values": self.values,
|
||||
"indices": self.indices,
|
||||
}
|
||||
|
||||
def as_dict(self) -> Dict[int, float]:
|
||||
return {i: v for i, v in zip(self.indices, self.values)}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[int, float]) -> "SparseEmbedding":
|
||||
indices, values = zip(*data.items())
|
||||
return cls(values=np.array(values), indices=np.array(indices))
|
||||
|
||||
|
||||
class SparseTextEmbeddingBase(ModelManagement):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.cache_dir = cache_dir
|
||||
self.threads = threads
|
||||
self._local_files_only = kwargs.pop("local_files_only", False)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[SparseEmbedding]:
|
||||
raise NotImplementedError()
|
||||
|
||||
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[SparseEmbedding]:
|
||||
"""
|
||||
Embeds a list of text passages into a list of embeddings.
|
||||
|
||||
Args:
|
||||
texts (Iterable[str]): The list of texts to embed.
|
||||
**kwargs: Additional keyword argument to pass to the embed method.
|
||||
|
||||
Yields:
|
||||
Iterable[SparseEmbedding]: The sparse embeddings.
|
||||
"""
|
||||
|
||||
# This is model-specific, so that different models can have specialized implementations
|
||||
yield from self.embed(texts, **kwargs)
|
||||
|
||||
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
|
||||
"""
|
||||
Embeds queries
|
||||
|
||||
Args:
|
||||
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
|
||||
|
||||
Returns:
|
||||
Iterable[SparseEmbedding]: The sparse embeddings.
|
||||
"""
|
||||
|
||||
# This is model-specific, so that different models can have specialized implementations
|
||||
if isinstance(query, str):
|
||||
yield from self.embed([query], **kwargs)
|
||||
if isinstance(query, Iterable):
|
||||
yield from self.embed(query, **kwargs)
|
||||
@@ -1,100 +0,0 @@
|
||||
from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.sparse.bm42 import Bm42
|
||||
from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
|
||||
from fastembed.sparse.splade_pp import SpladePP
|
||||
|
||||
|
||||
class SparseTextEmbedding(SparseTextEmbeddingBase):
|
||||
EMBEDDINGS_REGISTRY: List[Type[SparseTextEmbeddingBase]] = [
|
||||
SpladePP,
|
||||
Bm42,
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
|
||||
Example:
|
||||
```
|
||||
[
|
||||
{
|
||||
"model": "prithvida/SPLADE_PP_en_v1",
|
||||
"vocab_size": 30522,
|
||||
"description": "Independent Implementation of SPLADE++ Model for English",
|
||||
"size_in_GB": 0.532,
|
||||
"sources": {
|
||||
"hf": "qdrant/SPLADE_PP_en_v1",
|
||||
},
|
||||
}
|
||||
]
|
||||
```
|
||||
"""
|
||||
result = []
|
||||
for embedding in cls.EMBEDDINGS_REGISTRY:
|
||||
result.extend(embedding.list_supported_models())
|
||||
return result
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
|
||||
supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
|
||||
if any(model_name.lower() == model["model"].lower() for model in supported_models):
|
||||
self.model = EMBEDDING_MODEL_TYPE(
|
||||
model_name, cache_dir, threads, providers=providers, **kwargs
|
||||
)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
f"Model {model_name} is not supported in SparseTextEmbedding."
|
||||
"Please check the supported models using `SparseTextEmbedding.list_supported_models()`"
|
||||
)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[SparseEmbedding]:
|
||||
"""
|
||||
Encode a list of documents into list of embeddings.
|
||||
We use mean pooling with attention so that the model can handle variable-length inputs.
|
||||
|
||||
Args:
|
||||
documents: Iterator of documents or single document to embed
|
||||
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
|
||||
parallel:
|
||||
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
||||
If 0, use all available cores.
|
||||
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
||||
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
|
||||
|
||||
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
|
||||
"""
|
||||
Embeds queries
|
||||
|
||||
Args:
|
||||
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
|
||||
|
||||
Returns:
|
||||
Iterable[SparseEmbedding]: The sparse embeddings.
|
||||
"""
|
||||
yield from self.model.query_embed(query, **kwargs)
|
||||
@@ -1,139 +0,0 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.common.onnx_model import OnnxOutputContext
|
||||
from fastembed.common.utils import define_cache_dir
|
||||
from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
|
||||
from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
|
||||
|
||||
|
||||
supported_splade_models = [
|
||||
{
|
||||
"model": "prithvida/Splade_PP_en_v1",
|
||||
"vocab_size": 30522,
|
||||
"description": "Misspelled version of the model. Retained for backward compatibility. Independent Implementation of SPLADE++ Model for English",
|
||||
"size_in_GB": 0.532,
|
||||
"sources": {
|
||||
"hf": "Qdrant/SPLADE_PP_en_v1",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "prithivida/Splade_PP_en_v1",
|
||||
"vocab_size": 30522,
|
||||
"description": "Independent Implementation of SPLADE++ Model for English",
|
||||
"size_in_GB": 0.532,
|
||||
"sources": {
|
||||
"hf": "Qdrant/SPLADE_PP_en_v1",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
|
||||
def _post_process_onnx_output(
|
||||
self, output: OnnxOutputContext
|
||||
) -> Iterable[SparseEmbedding]:
|
||||
relu_log = np.log(1 + np.maximum(output.model_output, 0))
|
||||
|
||||
weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1)
|
||||
|
||||
scores = np.max(weighted_log, axis=1)
|
||||
|
||||
# Score matrix of shape (batch_size, vocab_size)
|
||||
# Most of the values are 0, only a few are non-zero
|
||||
for row_scores in scores:
|
||||
indices = row_scores.nonzero()[0]
|
||||
scores = row_scores[indices]
|
||||
yield SparseEmbedding(values=scores, indices=indices)
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
"""
|
||||
return supported_splade_models
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
model_name (str): The name of the model to use.
|
||||
cache_dir (str, optional): The path to the cache directory.
|
||||
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
||||
Defaults to `fastembed_cache` in the system's temp directory.
|
||||
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
|
||||
"""
|
||||
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
model_description = self._get_model_description(model_name)
|
||||
cache_dir = define_cache_dir(cache_dir)
|
||||
|
||||
model_dir = self.download_model(
|
||||
model_description, cache_dir, local_files_only=self._local_files_only
|
||||
)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
providers=providers,
|
||||
)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
**kwargs,
|
||||
) -> Iterable[SparseEmbedding]:
|
||||
"""
|
||||
Encode a list of documents into list of embeddings.
|
||||
We use mean pooling with attention so that the model can handle variable-length inputs.
|
||||
|
||||
Args:
|
||||
documents: Iterator of documents or single document to embed
|
||||
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
|
||||
parallel:
|
||||
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
||||
If 0, use all available cores.
|
||||
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
||||
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self._embed_documents(
|
||||
model_name=self.model_name,
|
||||
cache_dir=str(self.cache_dir),
|
||||
documents=documents,
|
||||
batch_size=batch_size,
|
||||
parallel=parallel,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
|
||||
return SpladePPEmbeddingWorker
|
||||
|
||||
|
||||
class SpladePPEmbeddingWorker(TextEmbeddingWorker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
) -> SpladePP:
|
||||
return SpladePP(model_name=model_name, cache_dir=cache_dir, threads=1)
|
||||
@@ -1,3 +0,0 @@
|
||||
from fastembed.text.text_embedding import TextEmbedding
|
||||
|
||||
__all__ = ["TextEmbedding"]
|
||||
|
||||
@@ -1,49 +0,0 @@
|
||||
from typing import Type, List, Dict, Any, Tuple, Iterable
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.onnx_model import OnnxOutputContext
|
||||
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
|
||||
from fastembed.text.onnx_text_model import TextEmbeddingWorker
|
||||
|
||||
supported_clip_models = [
|
||||
{
|
||||
"model": "Qdrant/clip-ViT-B-32-text",
|
||||
"dim": 512,
|
||||
"description": "CLIP text encoder",
|
||||
"size_in_GB": 0.25,
|
||||
"sources": {
|
||||
"hf": "Qdrant/clip-ViT-B-32-text",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class CLIPOnnxEmbedding(OnnxTextEmbedding):
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
|
||||
return CLIPEmbeddingWorker
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
"""
|
||||
return supported_clip_models
|
||||
|
||||
def _post_process_onnx_output(
|
||||
self, output: OnnxOutputContext
|
||||
) -> Iterable[np.ndarray]:
|
||||
return output.model_output
|
||||
|
||||
|
||||
class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
) -> OnnxTextEmbedding:
|
||||
return CLIPOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
|
||||
@@ -2,8 +2,7 @@ from typing import Type, List, Dict, Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
|
||||
from fastembed.text.onnx_text_model import TextEmbeddingWorker
|
||||
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker, EmbeddingWorker
|
||||
|
||||
supported_multilingual_e5_models = [
|
||||
{
|
||||
@@ -15,25 +14,22 @@ supported_multilingual_e5_models = [
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
|
||||
"hf": "qdrant/multilingual-e5-large-onnx",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
"additional_files": ["model.onnx_data"],
|
||||
},
|
||||
{
|
||||
"model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
|
||||
"dim": 768,
|
||||
"description": "Sentence-transformers model for tasks like clustering or semantic search",
|
||||
"size_in_GB": 1.00,
|
||||
"size_in_GB": 1.11,
|
||||
"sources": {
|
||||
"hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class E5OnnxEmbedding(OnnxTextEmbedding):
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
|
||||
def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
|
||||
return E5OnnxEmbeddingWorker
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -1,35 +1,31 @@
|
||||
from typing import Type, List, Dict, Any, Tuple, Iterable
|
||||
from typing import Type, List, Dict, Any, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.onnx_model import OnnxOutputContext
|
||||
from fastembed.common.utils import normalize
|
||||
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
|
||||
from fastembed.text.onnx_text_model import TextEmbeddingWorker
|
||||
from fastembed.common.models import normalize
|
||||
from fastembed.text.onnx_embedding import OnnxTextEmbedding, EmbeddingWorker, OnnxTextEmbeddingWorker
|
||||
|
||||
supported_jina_models = [
|
||||
{
|
||||
"model": "jinaai/jina-embeddings-v2-base-en",
|
||||
"dim": 768,
|
||||
"description": "English embedding model supporting 8192 sequence length",
|
||||
"size_in_GB": 0.52,
|
||||
"size_in_GB": 0.55,
|
||||
"sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "jinaai/jina-embeddings-v2-small-en",
|
||||
"dim": 512,
|
||||
"description": "English embedding model supporting 8192 sequence length",
|
||||
"size_in_GB": 0.12,
|
||||
"size_in_GB": 0.13,
|
||||
"sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class JinaOnnxEmbedding(OnnxTextEmbedding):
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
|
||||
def _get_worker_class(cls) -> Type[EmbeddingWorker]:
|
||||
return JinaEmbeddingWorker
|
||||
|
||||
@classmethod
|
||||
@@ -51,12 +47,10 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
|
||||
"""
|
||||
return supported_jina_models
|
||||
|
||||
def _post_process_onnx_output(
|
||||
self, output: OnnxOutputContext
|
||||
) -> Iterable[np.ndarray]:
|
||||
embeddings = output.model_output
|
||||
attn_mask = output.attention_mask
|
||||
return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
|
||||
@classmethod
|
||||
def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> np.ndarray:
|
||||
embeddings, attn_mask = output
|
||||
return normalize(cls.mean_pooling(embeddings, attn_mask)).astype(np.float32)
|
||||
|
||||
|
||||
class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):
|
||||
|
||||
+171
-127
@@ -1,77 +1,94 @@
|
||||
from typing import Dict, Optional, Union, Iterable, Type, List, Any, Sequence
|
||||
import os
|
||||
from multiprocessing import get_all_start_methods
|
||||
from typing import List, Dict, Any, Optional, Tuple, Union, Iterable, Type
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.common.onnx_model import OnnxOutputContext
|
||||
from fastembed.common.utils import normalize, define_cache_dir
|
||||
from fastembed.text.onnx_text_model import TextEmbeddingWorker, OnnxTextModel
|
||||
from fastembed.common.model_management import locate_model_file
|
||||
from fastembed.common.models import load_tokenizer, normalize
|
||||
from fastembed.common.utils import define_cache_dir, iter_batch
|
||||
from fastembed.parallel_processor import ParallelWorkerPool, Worker
|
||||
from fastembed.text.text_embedding_base import TextEmbeddingBase
|
||||
|
||||
|
||||
supported_onnx_models = [
|
||||
{
|
||||
"model": "BAAI/bge-base-en",
|
||||
"dim": 768,
|
||||
"description": "Base English model",
|
||||
"size_in_GB": 0.42,
|
||||
"size_in_GB": 0.5,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-base-en-v1.5",
|
||||
"dim": 768,
|
||||
"description": "Base English model, v1.5",
|
||||
"size_in_GB": 0.21,
|
||||
"size_in_GB": 0.44,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
|
||||
"hf": "qdrant/bge-base-en-v1.5-onnx-q",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-large-en-v1.5-quantized",
|
||||
"dim": 1024,
|
||||
"description": "Large English model, v1.5",
|
||||
"size_in_GB": 1.34,
|
||||
"sources": {
|
||||
"hf": "qdrant/bge-large-en-v1.5-onnx-q",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-large-en-v1.5",
|
||||
"dim": 1024,
|
||||
"description": "Large English model, v1.5",
|
||||
"size_in_GB": 1.20,
|
||||
"size_in_GB": 1.34,
|
||||
"sources": {
|
||||
"hf": "qdrant/bge-large-en-v1.5-onnx",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-small-en",
|
||||
"dim": 384,
|
||||
"description": "Fast English model",
|
||||
"size_in_GB": 0.13,
|
||||
"size_in_GB": 0.2,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
# {
|
||||
# "model": "BAAI/bge-small-en",
|
||||
# "dim": 384,
|
||||
# "description": "Fast English model",
|
||||
# "size_in_GB": 0.2,
|
||||
# "hf_sources": [],
|
||||
# "compressed_url_sources": [
|
||||
# "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en.tar.gz",
|
||||
# "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz"
|
||||
# ]
|
||||
# },
|
||||
{
|
||||
"model": "BAAI/bge-small-en-v1.5",
|
||||
"dim": 384,
|
||||
"description": "Fast and Default English model",
|
||||
"size_in_GB": 0.067,
|
||||
"size_in_GB": 0.13,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz",
|
||||
"hf": "qdrant/bge-small-en-v1.5-onnx-q",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-small-zh-v1.5",
|
||||
"dim": 512,
|
||||
"description": "Fast and recommended Chinese model",
|
||||
"size_in_GB": 0.09,
|
||||
"size_in_GB": 0.1,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
{ # todo: it is not a flag embedding
|
||||
"model": "sentence-transformers/all-MiniLM-L6-v2",
|
||||
"dim": 384,
|
||||
"description": "Sentence Transformer model, MiniLM-L6-v2",
|
||||
@@ -80,140 +97,96 @@ supported_onnx_models = [
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
|
||||
"hf": "qdrant/all-MiniLM-L6-v2-onnx",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
|
||||
"dim": 384,
|
||||
"description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
|
||||
"size_in_GB": 0.22,
|
||||
"size_in_GB": 0.46,
|
||||
"sources": {
|
||||
"hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "nomic-ai/nomic-embed-text-v1",
|
||||
"dim": 768,
|
||||
"description": "8192 context length english model",
|
||||
"size_in_GB": 0.52,
|
||||
"size_in_GB": 0.54,
|
||||
"sources": {
|
||||
"hf": "nomic-ai/nomic-embed-text-v1",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "nomic-ai/nomic-embed-text-v1.5",
|
||||
"dim": 768,
|
||||
"description": "8192 context length english model",
|
||||
"size_in_GB": 0.52,
|
||||
"size_in_GB": 0.54,
|
||||
"sources": {
|
||||
"hf": "nomic-ai/nomic-embed-text-v1.5",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "nomic-ai/nomic-embed-text-v1.5-Q",
|
||||
"dim": 768,
|
||||
"description": "Quantized 8192 context length english model",
|
||||
"size_in_GB": 0.13,
|
||||
"sources": {
|
||||
"hf": "nomic-ai/nomic-embed-text-v1.5",
|
||||
},
|
||||
"model_file": "onnx/model_quantized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "thenlper/gte-large",
|
||||
"dim": 1024,
|
||||
"description": "Large general text embeddings model",
|
||||
"size_in_GB": 1.20,
|
||||
"size_in_GB": 1.34,
|
||||
"sources": {
|
||||
"hf": "qdrant/gte-large-onnx",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "mixedbread-ai/mxbai-embed-large-v1",
|
||||
"dim": 1024,
|
||||
"description": "MixedBread Base sentence embedding model, does well on MTEB",
|
||||
"size_in_GB": 0.64,
|
||||
"sources": {
|
||||
"hf": "mixedbread-ai/mxbai-embed-large-v1",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "snowflake/snowflake-arctic-embed-xs",
|
||||
"dim": 384,
|
||||
"description": "Based on all-MiniLM-L6-v2 model with only 22m parameters, ideal for latency/TCO budgets.",
|
||||
"size_in_GB": 0.09,
|
||||
"sources": {
|
||||
"hf": "snowflake/snowflake-arctic-embed-xs",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "snowflake/snowflake-arctic-embed-s",
|
||||
"dim": 384,
|
||||
"description": "Based on infloat/e5-small-unsupervised, does not trade off retrieval accuracy for its small size.",
|
||||
"size_in_GB": 0.13,
|
||||
"sources": {
|
||||
"hf": "snowflake/snowflake-arctic-embed-s",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "snowflake/snowflake-arctic-embed-m",
|
||||
"dim": 768,
|
||||
"description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
|
||||
"size_in_GB": 0.43,
|
||||
"sources": {
|
||||
"hf": "Snowflake/snowflake-arctic-embed-m",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "snowflake/snowflake-arctic-embed-m-long",
|
||||
"dim": 768,
|
||||
"description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
|
||||
"size_in_GB": 0.54,
|
||||
"sources": {
|
||||
"hf": "snowflake/snowflake-arctic-embed-m-long",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "snowflake/snowflake-arctic-embed-l",
|
||||
"dim": 1024,
|
||||
"description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
|
||||
"size_in_GB": 1.02,
|
||||
"sources": {
|
||||
"hf": "snowflake/snowflake-arctic-embed-l",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
# {
|
||||
# "model": "sentence-transformers/all-MiniLM-L6-v2",
|
||||
# "dim": 384,
|
||||
# "description": "Sentence Transformer model, MiniLM-L6-v2",
|
||||
# "size_in_GB": 0.09,
|
||||
# "hf_sources": [
|
||||
# "qdrant/all-MiniLM-L6-v2-onnx"
|
||||
# ],
|
||||
# "compressed_url_sources": [
|
||||
# "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",
|
||||
# "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
|
||||
# ]
|
||||
# }
|
||||
]
|
||||
|
||||
|
||||
class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
|
||||
class OnnxTextEmbedding(TextEmbeddingBase):
|
||||
"""Implementation of the Flag Embedding model."""
|
||||
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Lists the supported models.
|
||||
"""Lists the supported models.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries containing the model information.
|
||||
"""
|
||||
return supported_onnx_models
|
||||
|
||||
@classmethod
|
||||
def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Gets the model description from the model_name.
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model.
|
||||
|
||||
raises:
|
||||
ValueError: If the model_name is not supported.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: The model description.
|
||||
"""
|
||||
for model in cls.list_supported_models():
|
||||
if model_name == model["model"]:
|
||||
return model
|
||||
|
||||
raise ValueError(f"Model {model_name} is not supported in FlagEmbedding.")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "BAAI/bge-small-en-v1.5",
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
@@ -230,18 +203,27 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
|
||||
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
model_description = self._get_model_description(model_name)
|
||||
cache_dir = define_cache_dir(cache_dir)
|
||||
model_dir = self.download_model(
|
||||
model_description, cache_dir, local_files_only=self._local_files_only
|
||||
)
|
||||
self.model_name = model_name
|
||||
self._model_description = self._get_model_description(model_name)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
providers=providers,
|
||||
)
|
||||
self._cache_dir = define_cache_dir(cache_dir)
|
||||
self._model_dir = self.download_model(self._model_description, self._cache_dir)
|
||||
self._max_length = 512
|
||||
|
||||
model_path = locate_model_file(self._model_dir, ["model.onnx", "model_optimized.onnx"])
|
||||
|
||||
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
|
||||
onnx_providers = ["CPUExecutionProvider"]
|
||||
|
||||
so = ort.SessionOptions()
|
||||
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
if self.threads is not None:
|
||||
so.intra_op_num_threads = self.threads
|
||||
so.inter_op_num_threads = self.threads
|
||||
|
||||
self.tokenizer = load_tokenizer(model_dir=self._model_dir, max_length=self._max_length)
|
||||
self.model = ort.InferenceSession(str(model_path), providers=onnx_providers, sess_options=so)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
@@ -265,32 +247,94 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
|
||||
Returns:
|
||||
List of embeddings, one per document
|
||||
"""
|
||||
yield from self._embed_documents(
|
||||
model_name=self.model_name,
|
||||
cache_dir=str(self.cache_dir),
|
||||
documents=documents,
|
||||
batch_size=batch_size,
|
||||
parallel=parallel,
|
||||
)
|
||||
is_small = False
|
||||
|
||||
if isinstance(documents, str):
|
||||
documents = [documents]
|
||||
is_small = True
|
||||
|
||||
if isinstance(documents, list):
|
||||
if len(documents) < batch_size:
|
||||
is_small = True
|
||||
|
||||
if parallel == 0:
|
||||
parallel = os.cpu_count()
|
||||
|
||||
if parallel is None or is_small:
|
||||
for batch in iter_batch(documents, batch_size):
|
||||
yield from self._post_process_onnx_output(self.onnx_embed(batch))
|
||||
else:
|
||||
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
|
||||
params = {
|
||||
"model_name": self.model_name,
|
||||
"cache_dir": str(self._cache_dir),
|
||||
}
|
||||
pool = ParallelWorkerPool(parallel, self._get_worker_class(), start_method=start_method)
|
||||
for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
|
||||
yield from self._post_process_onnx_output(batch)
|
||||
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
|
||||
def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
|
||||
return OnnxTextEmbeddingWorker
|
||||
|
||||
def _preprocess_onnx_input(
|
||||
self, onnx_input: Dict[str, np.ndarray], **kwargs
|
||||
) -> Dict[str, np.ndarray]:
|
||||
def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Preprocess the onnx input.
|
||||
"""
|
||||
return onnx_input
|
||||
|
||||
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
|
||||
embeddings = output.model_output
|
||||
@classmethod
|
||||
def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]):
|
||||
embeddings, _ = output
|
||||
return normalize(embeddings[:, 0]).astype(np.float32)
|
||||
|
||||
def onnx_embed(self, documents: List[str]) -> Tuple[np.ndarray, np.ndarray]:
|
||||
encoded = self.tokenizer.encode_batch(documents)
|
||||
input_ids = np.array([e.ids for e in encoded])
|
||||
attention_mask = np.array([e.attention_mask for e in encoded])
|
||||
|
||||
class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
|
||||
onnx_input = {
|
||||
"input_ids": np.array(input_ids, dtype=np.int64),
|
||||
"attention_mask": np.array(attention_mask, dtype=np.int64),
|
||||
"token_type_ids": np.array([np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64),
|
||||
}
|
||||
|
||||
onnx_input = self._preprocess_onnx_input(onnx_input)
|
||||
|
||||
model_output = self.model.run(None, onnx_input)
|
||||
embeddings = model_output[0]
|
||||
return embeddings, attention_mask
|
||||
|
||||
|
||||
class EmbeddingWorker(Worker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
) -> OnnxTextEmbedding:
|
||||
raise NotImplementedError()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
):
|
||||
self.model = self.init_embedding(model_name, cache_dir)
|
||||
|
||||
@classmethod
|
||||
def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
|
||||
return cls(
|
||||
model_name=model_name,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
|
||||
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
|
||||
for idx, batch in items:
|
||||
embeddings, attn_mask = self.model.onnx_embed(batch)
|
||||
yield idx, (embeddings, attn_mask)
|
||||
|
||||
|
||||
class OnnxTextEmbeddingWorker(EmbeddingWorker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
|
||||
@@ -1,123 +0,0 @@
|
||||
import os
|
||||
from multiprocessing import get_all_start_methods
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union, Sequence
|
||||
|
||||
import numpy as np
|
||||
from tokenizers import Encoding
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.common.preprocessor_utils import load_tokenizer
|
||||
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
|
||||
from fastembed.common.utils import iter_batch
|
||||
from fastembed.parallel_processor import ParallelWorkerPool
|
||||
|
||||
|
||||
class OnnxTextModel(OnnxModel[T]):
|
||||
ONNX_OUTPUT_NAMES: Optional[List[str]] = None
|
||||
|
||||
@classmethod
|
||||
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
|
||||
raise NotImplementedError("Subclasses must implement this method")
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.tokenizer = None
|
||||
self.special_token_to_id = {}
|
||||
|
||||
def _preprocess_onnx_input(
|
||||
self, onnx_input: Dict[str, np.ndarray], **kwargs
|
||||
) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Preprocess the onnx input.
|
||||
"""
|
||||
return onnx_input
|
||||
|
||||
def load_onnx_model(
|
||||
self,
|
||||
model_dir: Path,
|
||||
model_file: str,
|
||||
threads: Optional[int],
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
) -> None:
|
||||
super().load_onnx_model(
|
||||
model_dir=model_dir, model_file=model_file, threads=threads, providers=providers
|
||||
)
|
||||
self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
|
||||
|
||||
def tokenize(self, documents: List[str], **kwargs) -> List[Encoding]:
|
||||
return self.tokenizer.encode_batch(documents)
|
||||
|
||||
def onnx_embed(
|
||||
self,
|
||||
documents: List[str],
|
||||
**kwargs,
|
||||
) -> OnnxOutputContext:
|
||||
encoded = self.tokenize(documents, **kwargs)
|
||||
input_ids = np.array([e.ids for e in encoded])
|
||||
attention_mask = np.array([e.attention_mask for e in encoded])
|
||||
input_names = {node.name for node in self.model.get_inputs()}
|
||||
onnx_input = {
|
||||
"input_ids": np.array(input_ids, dtype=np.int64),
|
||||
}
|
||||
if "attention_mask" in input_names:
|
||||
onnx_input["attention_mask"] = np.array(attention_mask, dtype=np.int64)
|
||||
if "token_type_ids" in input_names:
|
||||
onnx_input["token_type_ids"] = np.array(
|
||||
[np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
|
||||
)
|
||||
|
||||
onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
|
||||
|
||||
model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
|
||||
return OnnxOutputContext(
|
||||
model_output=model_output[0],
|
||||
attention_mask=onnx_input.get("attention_mask", attention_mask),
|
||||
input_ids=onnx_input.get("input_ids", input_ids),
|
||||
)
|
||||
|
||||
def _embed_documents(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
) -> Iterable[T]:
|
||||
is_small = False
|
||||
|
||||
if isinstance(documents, str):
|
||||
documents = [documents]
|
||||
is_small = True
|
||||
|
||||
if isinstance(documents, list):
|
||||
if len(documents) < batch_size:
|
||||
is_small = True
|
||||
|
||||
if parallel == 0:
|
||||
parallel = os.cpu_count()
|
||||
|
||||
if parallel is None or is_small:
|
||||
for batch in iter_batch(documents, batch_size):
|
||||
yield from self._post_process_onnx_output(self.onnx_embed(batch))
|
||||
else:
|
||||
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
|
||||
params = {
|
||||
"model_name": model_name,
|
||||
"cache_dir": cache_dir,
|
||||
}
|
||||
pool = ParallelWorkerPool(
|
||||
parallel, self._get_worker_class(), start_method=start_method
|
||||
)
|
||||
for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
|
||||
yield from self._post_process_onnx_output(batch)
|
||||
|
||||
|
||||
class TextEmbeddingWorker(EmbeddingWorker):
|
||||
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
|
||||
for idx, batch in items:
|
||||
onnx_output = self.model.onnx_embed(batch)
|
||||
yield idx, onnx_output
|
||||
@@ -1,9 +1,7 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional, Type, Union, Sequence
|
||||
from typing import Optional, Union, Iterable, List, Dict, Any, Type
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.text.clip_embedding import CLIPOnnxEmbedding
|
||||
from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
|
||||
from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
|
||||
from fastembed.text.onnx_embedding import OnnxTextEmbedding
|
||||
@@ -15,7 +13,6 @@ class TextEmbedding(TextEmbeddingBase):
|
||||
OnnxTextEmbedding,
|
||||
E5OnnxEmbedding,
|
||||
JinaOnnxEmbedding,
|
||||
CLIPOnnxEmbedding,
|
||||
]
|
||||
|
||||
@classmethod
|
||||
@@ -52,17 +49,14 @@ class TextEmbedding(TextEmbeddingBase):
|
||||
model_name: str = "BAAI/bge-small-en-v1.5",
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
|
||||
supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
|
||||
if any(model_name.lower() == model["model"].lower() for model in supported_models):
|
||||
self.model = EMBEDDING_MODEL_TYPE(
|
||||
model_name, cache_dir, threads, providers=providers, **kwargs
|
||||
)
|
||||
for embedding in self.EMBEDDINGS_REGISTRY:
|
||||
supported_models = embedding.list_supported_models()
|
||||
if any(model_name == model["model"] for model in supported_models):
|
||||
self.model = embedding(model_name, cache_dir, threads, **kwargs)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from typing import Iterable, Optional, Union
|
||||
from typing import Any, Dict, Iterable, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -6,17 +6,14 @@ from fastembed.common.model_management import ModelManagement
|
||||
|
||||
|
||||
class TextEmbeddingBase(ModelManagement):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
@classmethod
|
||||
def list_supported_models(cls) -> List[Dict[str, Any]]:
|
||||
raise NotImplementedError()
|
||||
|
||||
def __init__(self, model_name: str, cache_dir: Optional[str] = None, threads: Optional[int] = None, **kwargs):
|
||||
self.model_name = model_name
|
||||
self.cache_dir = cache_dir
|
||||
self.threads = threads
|
||||
self._local_files_only = kwargs.pop("local_files_only", False)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
|
||||
Generated
+3433
File diff suppressed because it is too large
Load Diff
+11
-15
@@ -1,8 +1,8 @@
|
||||
[tool.poetry]
|
||||
name = "fastembed-gpu"
|
||||
version = "0.3.0"
|
||||
name = "fastembed"
|
||||
version = "0.2.2"
|
||||
description = "Fast, light, accurate library built for retrieval embedding generation"
|
||||
authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
|
||||
authors = ["NirantK <nirant.bits@gmail.com>"]
|
||||
license = "Apache License"
|
||||
readme = "README.md"
|
||||
packages = [{include = "fastembed"}]
|
||||
@@ -12,33 +12,29 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.8.0,<3.13"
|
||||
onnxruntime-gpu = "^1.17.0"
|
||||
onnx = "^1.15.0"
|
||||
onnxruntime = "^1.17.0"
|
||||
tqdm = "^4.66"
|
||||
requests = "^2.31"
|
||||
tokenizers = ">=0.15,<1.0"
|
||||
huggingface-hub = ">=0.20,<1.0"
|
||||
tokenizers = "^0.15.1"
|
||||
huggingface-hub = "^0.20"
|
||||
loguru = "^0.7.2"
|
||||
numpy = [
|
||||
{ version = ">=1.21", python = "<3.12" },
|
||||
{ version = ">=1.26", python = ">=3.12" }
|
||||
]
|
||||
pillow = "^10.3.0"
|
||||
snowballstemmer = "^2.2.0"
|
||||
PyStemmer = "^2.2.0"
|
||||
mmh3 = "^4.0"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
pytest = "^7.4.2"
|
||||
ruff = ">=0.3.1,<1.0"
|
||||
ruff = "^0.2.2"
|
||||
notebook = ">=7.0.2"
|
||||
pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
|
||||
|
||||
[tool.poetry.group.docs.dependencies]
|
||||
mkdocs-material = "^9.5.10"
|
||||
mkdocstrings = "^0.24.0"
|
||||
pillow = "^10.2.0"
|
||||
cairosvg = "^2.7.1"
|
||||
mknotebooks = "^0.8.0"
|
||||
pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
|
||||
click = "^8.1.7"
|
||||
|
||||
|
||||
[build-system]
|
||||
@@ -46,4 +42,4 @@ requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 99
|
||||
line-length = 120
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
from pathlib import Path
|
||||
|
||||
TEST_DIR = Path(__file__).parent
|
||||
TEST_MISC_DIR = TEST_DIR / "misc"
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 169 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 11 KiB |
@@ -1,58 +0,0 @@
|
||||
import numpy as np
|
||||
|
||||
from fastembed import SparseTextEmbedding
|
||||
|
||||
|
||||
def test_attention_embeddings():
|
||||
model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
|
||||
|
||||
output = list(
|
||||
model.query_embed(
|
||||
[
|
||||
"I must not fear. Fear is the mind-killer.",
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(output) == 1
|
||||
|
||||
for result in output:
|
||||
assert len(result.indices) == len(result.values)
|
||||
assert np.allclose(result.values, np.ones(len(result.values)))
|
||||
|
||||
quotes = [
|
||||
"I must not fear. Fear is the mind-killer.",
|
||||
"All animals are equal, but some animals are more equal than others.",
|
||||
"It was a pleasure to burn.",
|
||||
"The sky above the port was the color of television, tuned to a dead channel.",
|
||||
"In the beginning, the universe was created."
|
||||
" This has made a lot of people very angry and been widely regarded as a bad move.",
|
||||
"It's a truth universally acknowledged that a zombie in possession of brains must be in want of more brains.",
|
||||
"War is peace. Freedom is slavery. Ignorance is strength.",
|
||||
"We're not in Infinity; we're in the suburbs.",
|
||||
"I was a thousand times more evil than thou!",
|
||||
"History is merely a list of surprises... It can only prepare us to be surprised yet again.",
|
||||
]
|
||||
|
||||
output = list(model.embed(quotes))
|
||||
|
||||
assert len(output) == len(quotes)
|
||||
|
||||
for result in output:
|
||||
assert len(result.indices) == len(result.values)
|
||||
assert len(result.indices) > 0
|
||||
|
||||
# Test support for unknown languages
|
||||
output = list(
|
||||
model.query_embed(
|
||||
[
|
||||
"привет мир!",
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert len(output) == 1
|
||||
|
||||
for result in output:
|
||||
assert len(result.indices) == len(result.values)
|
||||
assert len(result.indices) == 2
|
||||
@@ -1,73 +0,0 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from fastembed import ImageEmbedding
|
||||
from tests.config import TEST_MISC_DIR
|
||||
|
||||
CANONICAL_VECTOR_VALUES = {
|
||||
"Qdrant/clip-ViT-B-32-vision": np.array([-0.0098, 0.0128, -0.0274, 0.002, -0.0059]),
|
||||
"Qdrant/resnet50-onnx": np.array(
|
||||
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01046245, 0.01171397, 0.00705971, 0.0]
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def test_embedding():
|
||||
is_ci = os.getenv("CI")
|
||||
|
||||
for model_desc in ImageEmbedding.list_supported_models():
|
||||
if not is_ci and model_desc["size_in_GB"] > 1:
|
||||
continue
|
||||
|
||||
dim = model_desc["dim"]
|
||||
|
||||
model = ImageEmbedding(model_name=model_desc["model"])
|
||||
|
||||
images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")]
|
||||
embeddings = list(model.embed(images))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
assert embeddings.shape == (2, dim)
|
||||
|
||||
canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
|
||||
|
||||
assert np.allclose(
|
||||
embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
|
||||
), model_desc["model"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
|
||||
def test_batch_embedding(n_dims, model_name):
|
||||
model = ImageEmbedding(model_name=model_name)
|
||||
n_images = 32
|
||||
images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
|
||||
n_images // 2
|
||||
)
|
||||
|
||||
embeddings = list(model.embed(images, batch_size=10))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
|
||||
assert embeddings.shape == (n_images, n_dims)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
|
||||
def test_parallel_processing(n_dims, model_name):
|
||||
model = ImageEmbedding(model_name=model_name)
|
||||
|
||||
n_images = 32
|
||||
images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
|
||||
n_images // 2
|
||||
)
|
||||
embeddings = list(model.embed(images, batch_size=10, parallel=2))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
|
||||
embeddings_2 = list(model.embed(images, batch_size=10, parallel=None))
|
||||
embeddings_2 = np.stack(embeddings_2, axis=0)
|
||||
|
||||
embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
|
||||
embeddings_3 = np.stack(embeddings_3, axis=0)
|
||||
|
||||
assert embeddings.shape == (n_images, n_dims)
|
||||
assert np.allclose(embeddings, embeddings_2, atol=1e-3)
|
||||
assert np.allclose(embeddings, embeddings_3, atol=1e-3)
|
||||
@@ -1,113 +0,0 @@
|
||||
import numpy as np
|
||||
|
||||
from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding
|
||||
|
||||
|
||||
# vectors are abridged and rounded for brevity
|
||||
CANONICAL_COLUMN_VALUES = {
|
||||
"colbert-ir/colbertv2.0": np.array(
|
||||
[
|
||||
[0.0759, 0.0841, -0.0299, 0.0374, 0.0254],
|
||||
[0.0005, -0.0163, -0.0127, 0.2165, 0.1517],
|
||||
[-0.0257, -0.0575, 0.0135, 0.2202, 0.1896],
|
||||
[0.0846, 0.0122, 0.0032, -0.0109, -0.1041],
|
||||
[0.0477, 0.1078, -0.0314, 0.016, 0.0156],
|
||||
]
|
||||
)
|
||||
}
|
||||
|
||||
CANONICAL_QUERY_VALUES = {
|
||||
"colbert-ir/colbertv2.0": np.array(
|
||||
[
|
||||
[0.0824, 0.0872, -0.0324, 0.0418, 0.024],
|
||||
[-0.0007, -0.0154, -0.0113, 0.2277, 0.1528],
|
||||
[-0.0251, -0.0565, 0.0136, 0.2236, 0.1838],
|
||||
[0.0848, 0.0056, 0.0041, -0.0036, -0.1032],
|
||||
[0.0574, 0.1072, -0.0332, 0.0233, 0.0209],
|
||||
[0.1041, 0.0364, -0.0058, -0.027, -0.0704],
|
||||
[0.106, 0.0371, -0.0055, -0.0339, -0.0719],
|
||||
[0.1063, 0.0363, 0.0014, -0.0334, -0.0698],
|
||||
[0.112, 0.036, 0.0026, -0.0355, -0.0675],
|
||||
[0.1184, 0.0441, 0.0166, -0.0169, -0.0244],
|
||||
[0.1033, 0.035, 0.0183, 0.0475, 0.0612],
|
||||
[-0.0028, -0.014, -0.016, 0.2175, 0.1537],
|
||||
[0.0547, 0.0219, -0.007, 0.1748, 0.1154],
|
||||
[-0.001, -0.0184, -0.0112, 0.2197, 0.1523],
|
||||
[-0.0012, -0.0149, -0.0119, 0.2147, 0.152],
|
||||
[-0.0186, -0.0239, -0.014, 0.2196, 0.156],
|
||||
[-0.017, -0.0232, -0.0108, 0.2212, 0.157],
|
||||
[-0.0109, -0.0024, -0.003, 0.1972, 0.1391],
|
||||
[0.0898, 0.0219, -0.0255, 0.0734, -0.0096],
|
||||
[0.1143, 0.015, -0.022, 0.0417, -0.0421],
|
||||
[0.1056, 0.0091, -0.0137, 0.0129, -0.0619],
|
||||
[0.0234, 0.004, -0.0285, 0.1565, 0.0883],
|
||||
[-0.0037, -0.0079, -0.0204, 0.1982, 0.1502],
|
||||
[0.0988, 0.0377, 0.0226, 0.0309, 0.0508],
|
||||
[-0.0103, -0.0128, -0.0035, 0.2114, 0.155],
|
||||
[-0.0103, -0.0184, -0.011, 0.2252, 0.157],
|
||||
[-0.0033, -0.0292, -0.0097, 0.2237, 0.1607],
|
||||
[-0.0198, -0.0257, -0.0193, 0.2265, 0.165],
|
||||
[-0.0227, -0.0028, -0.0084, 0.1995, 0.1306],
|
||||
[0.0916, 0.0185, -0.0186, 0.0173, -0.0577],
|
||||
[0.1022, 0.0228, -0.0174, -0.0102, -0.065],
|
||||
[0.1043, 0.0231, -0.0144, -0.0246, -0.067],
|
||||
]
|
||||
)
|
||||
}
|
||||
|
||||
docs = ["Hello World"]
|
||||
|
||||
|
||||
def test_batch_embedding():
|
||||
docs_to_embed = docs * 10
|
||||
|
||||
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
|
||||
print("evaluating", model_name)
|
||||
model = LateInteractionTextEmbedding(model_name=model_name)
|
||||
result = list(model.embed(docs_to_embed, batch_size=6))
|
||||
|
||||
for value in result:
|
||||
token_num, abridged_dim = expected_result.shape
|
||||
assert np.allclose(value[:, :abridged_dim], expected_result, atol=10e-4)
|
||||
|
||||
|
||||
def test_single_embedding():
|
||||
docs_to_embed = docs
|
||||
|
||||
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
|
||||
print("evaluating", model_name)
|
||||
model = LateInteractionTextEmbedding(model_name=model_name, cache_dir="colbert-cache")
|
||||
result = next(iter(model.embed(docs_to_embed, batch_size=6)))
|
||||
token_num, abridged_dim = expected_result.shape
|
||||
assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4)
|
||||
|
||||
|
||||
def test_single_embedding_query():
|
||||
queries_to_embed = docs
|
||||
|
||||
for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
|
||||
print("evaluating", model_name)
|
||||
model = LateInteractionTextEmbedding(model_name=model_name, cache_dir="colbert-cache")
|
||||
result = next(iter(model.query_embed(queries_to_embed)))
|
||||
token_num, abridged_dim = expected_result.shape
|
||||
assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4)
|
||||
|
||||
|
||||
def test_parallel_processing():
|
||||
model = LateInteractionTextEmbedding(
|
||||
model_name="colbert-ir/colbertv2.0", cache_dir="colbert-cache"
|
||||
)
|
||||
token_dim = 128
|
||||
docs = ["hello world", "flag embedding"] * 100
|
||||
embeddings = list(model.embed(docs, batch_size=10, parallel=2))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
|
||||
embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
|
||||
embeddings_2 = np.stack(embeddings_2, axis=0)
|
||||
|
||||
embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
|
||||
embeddings_3 = np.stack(embeddings_3, axis=0)
|
||||
|
||||
assert embeddings.shape[0] == len(docs) and embeddings.shape[-1] == token_dim
|
||||
assert np.allclose(embeddings, embeddings_2, atol=1e-3)
|
||||
assert np.allclose(embeddings, embeddings_3, atol=1e-3)
|
||||
@@ -1,97 +0,0 @@
|
||||
import pytest
|
||||
|
||||
from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
|
||||
|
||||
CANONICAL_COLUMN_VALUES = {
|
||||
"prithvida/Splade_PP_en_v1": {
|
||||
"indices": [
|
||||
2040,
|
||||
2047,
|
||||
2088,
|
||||
2299,
|
||||
2748,
|
||||
3011,
|
||||
3376,
|
||||
3795,
|
||||
4774,
|
||||
5304,
|
||||
5798,
|
||||
6160,
|
||||
7592,
|
||||
7632,
|
||||
8484,
|
||||
],
|
||||
"values": [
|
||||
0.4219532012939453,
|
||||
0.4320072531700134,
|
||||
2.766580104827881,
|
||||
0.3314574658870697,
|
||||
1.395172119140625,
|
||||
0.021595917642116547,
|
||||
0.43770670890808105,
|
||||
0.0008370947907678783,
|
||||
0.5187209844589233,
|
||||
0.17124654352664948,
|
||||
0.14742016792297363,
|
||||
0.8142819404602051,
|
||||
2.803262710571289,
|
||||
2.1904349327087402,
|
||||
1.0531445741653442,
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
docs = ["Hello World"]
|
||||
|
||||
|
||||
def test_batch_embedding():
|
||||
docs_to_embed = docs * 10
|
||||
|
||||
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
|
||||
model = SparseTextEmbedding(model_name=model_name)
|
||||
result = next(iter(model.embed(docs_to_embed, batch_size=6)))
|
||||
assert result.indices.tolist() == expected_result["indices"]
|
||||
|
||||
for i, value in enumerate(result.values):
|
||||
assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
|
||||
|
||||
|
||||
def test_single_embedding():
|
||||
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
|
||||
model = SparseTextEmbedding(model_name=model_name)
|
||||
|
||||
passage_result = next(iter(model.embed(docs, batch_size=6)))
|
||||
query_result = next(iter(model.query_embed(docs)))
|
||||
for result in [passage_result, query_result]:
|
||||
assert result.indices.tolist() == expected_result["indices"]
|
||||
|
||||
for i, value in enumerate(result.values):
|
||||
assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
|
||||
|
||||
|
||||
def test_parallel_processing():
|
||||
import numpy as np
|
||||
|
||||
model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
|
||||
docs = ["hello world", "flag embedding"] * 30
|
||||
sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
|
||||
sparse_embeddings_all = list(model.embed(docs, batch_size=10, parallel=0))
|
||||
sparse_embeddings = list(model.embed(docs, batch_size=10, parallel=None))
|
||||
|
||||
assert (
|
||||
len(sparse_embeddings)
|
||||
== len(sparse_embeddings_duo)
|
||||
== len(sparse_embeddings_all)
|
||||
== len(docs)
|
||||
)
|
||||
|
||||
for sparse_embedding, sparse_embedding_duo, sparse_embedding_all in zip(
|
||||
sparse_embeddings, sparse_embeddings_duo, sparse_embeddings_all
|
||||
):
|
||||
assert (
|
||||
sparse_embedding.indices.tolist()
|
||||
== sparse_embedding_duo.indices.tolist()
|
||||
== sparse_embedding_all.indices.tolist()
|
||||
)
|
||||
assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3)
|
||||
assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)
|
||||
@@ -7,29 +7,14 @@ from fastembed.text.text_embedding import TextEmbedding
|
||||
|
||||
CANONICAL_VECTOR_VALUES = {
|
||||
"BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
|
||||
"BAAI/bge-small-en-v1.5": np.array(
|
||||
[0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]
|
||||
),
|
||||
"BAAI/bge-small-en-v1.5-quantized": np.array(
|
||||
[0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]
|
||||
),
|
||||
"BAAI/bge-small-zh-v1.5": np.array(
|
||||
[-0.01023294, 0.07634465, 0.0691722, -0.04458365, -0.03160762]
|
||||
),
|
||||
"BAAI/bge-small-en-v1.5": np.array([0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]),
|
||||
"BAAI/bge-small-zh-v1.5": np.array([-0.01023294, 0.07634465, 0.0691722, -0.04458365, -0.03160762]),
|
||||
"BAAI/bge-base-en": np.array([0.0115, 0.0372, 0.0295, 0.0121, 0.0346]),
|
||||
"BAAI/bge-base-en-v1.5": np.array(
|
||||
[0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]
|
||||
),
|
||||
"BAAI/bge-large-en-v1.5": np.array(
|
||||
[0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]
|
||||
),
|
||||
"BAAI/bge-large-en-v1.5-quantized": np.array(
|
||||
[0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]
|
||||
),
|
||||
"BAAI/bge-base-en-v1.5": np.array([0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]),
|
||||
"BAAI/bge-large-en-v1.5": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
|
||||
"BAAI/bge-large-en-v1.5-quantized": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
|
||||
"sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
|
||||
"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array(
|
||||
[0.0094, 0.0184, 0.0328, 0.0072, -0.0351]
|
||||
),
|
||||
"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array([0.0094, 0.0184, 0.0328, 0.0072, -0.0351]),
|
||||
"intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
|
||||
"sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
|
||||
[-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
|
||||
@@ -40,35 +25,18 @@ CANONICAL_VECTOR_VALUES = {
|
||||
"nomic-ai/nomic-embed-text-v1.5": np.array(
|
||||
[-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
|
||||
),
|
||||
"nomic-ai/nomic-embed-text-v1.5-Q": np.array(
|
||||
[-0.01554983, 0.0129992, -0.17909265, -0.01062993, 0.00512859]
|
||||
),
|
||||
"thenlper/gte-large": np.array(
|
||||
[-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
|
||||
),
|
||||
"mixedbread-ai/mxbai-embed-large-v1": np.array(
|
||||
[0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]
|
||||
),
|
||||
"snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]),
|
||||
"snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]),
|
||||
"snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]),
|
||||
"snowflake/snowflake-arctic-embed-m-long": np.array(
|
||||
[0.0080, -0.0266, -0.0335, 0.0282, 0.0143]
|
||||
),
|
||||
"snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]),
|
||||
"Qdrant/clip-ViT-B-32-text": np.array([0.0083, 0.0103, -0.0138, 0.0199, -0.0069]),
|
||||
"thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
|
||||
}
|
||||
|
||||
|
||||
def test_embedding():
|
||||
is_ci = os.getenv("CI")
|
||||
is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
|
||||
|
||||
for model_desc in TextEmbedding.list_supported_models():
|
||||
if not is_ci and model_desc["size_in_GB"] > 1:
|
||||
if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
|
||||
continue
|
||||
|
||||
dim = model_desc["dim"]
|
||||
|
||||
model = TextEmbedding(model_name=model_desc["model"])
|
||||
|
||||
docs = ["hello world", "flag embedding"]
|
||||
@@ -77,14 +45,11 @@ def test_embedding():
|
||||
assert embeddings.shape == (2, dim)
|
||||
|
||||
canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
|
||||
assert np.allclose(
|
||||
embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
|
||||
), model_desc["model"]
|
||||
assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"n_dims,model_name",
|
||||
[(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
|
||||
"n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")]
|
||||
)
|
||||
def test_batch_embedding(n_dims, model_name):
|
||||
model = TextEmbedding(model_name=model_name)
|
||||
@@ -97,8 +62,7 @@ def test_batch_embedding(n_dims, model_name):
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"n_dims,model_name",
|
||||
[(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
|
||||
"n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")]
|
||||
)
|
||||
def test_parallel_processing(n_dims, model_name):
|
||||
model = TextEmbedding(model_name=model_name)
|
||||
|
||||
Reference in New Issue
Block a user