diff --git a/README.md b/README.md index 2222255..e12aa2d 100644 --- a/README.md +++ b/README.md @@ -16,14 +16,30 @@ -NIFE compresses large embedding models into static, drop-in replacements with up to 200x faster query embedding [see benchmarks](). + +
+

+ Models | + Datasets | + Benchmarks | + Create your own model +

+ +NIFE compresses large embedding models into static, drop-in replacements with up to 200x faster query embedding [see benchmarks](#benchmarks). ## Features -- 200x faster CPU query embedding +- 400-900x faster CPU query embedding - Fully aligned with their teacher models - Re-use your existing vector index +## Table of contents + +1. [Quickstart](#quickstart) +2. [Installation](#installation) +3. [Usage](#usage) +4. [Rationale](#rationale) + ## Introduction Nearly Inference Free Embedding (NIFE) models are [static embedding](https://huggingface.co/blog/static-embeddings) models that are fully aligned with a much larger model. Because static models are so small and fast, NIFE allows you to: @@ -113,6 +129,13 @@ print(model.similarity(query, docs)) ``` +## Pretrained models + +I have two pretrained models: + +* [`stephantulkens/NIFE-mxbai-embed-large-v1`](https://huggingface.co/stephantulkens/NIFE-mxbai-embed-large-v1): aligned with [`mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1). +* [`stephantulkens/NIFE-gte-modernbert-base`](https://huggingface.co/stephantulkens/NIFE-gte-modernbert-base): aligned with [`gte-modernbert-base`](https://huggingface.co/Alibaba-NLP/gte-modernbert-base). + ## Rationale For retrieval using dense models, the normal mode of operation is to embed your documents, and put them in some index. Then, using that same model, also embed your queries. In general, larger embedding models are better than smaller models, so you're often better off by making your embedder as large as possible. This however, makes inference more difficult; you need to host a larger model, and embedding queries might take longer. @@ -130,6 +153,31 @@ See this table: As in doc-SPLADE, you lose performance. No real way about it, but as with other fast models, the gap is smaller than you might think. +## Benchmarks + +I benchmark our models on [NanoBEIR](https://huggingface.co/collections/zeta-alpha-ai/nanobeir). I use two trained models: + +* [`stephantulkens/NIFE-mxbai-embed-large-v1`](https://huggingface.co/stephantulkens/NIFE-mxbai-embed-large-v1) +* [`stephantulkens/NIFE-gte-modernbert-base`](https://huggingface.co/stephantulkens/NIFE-gte-modernbert-base) + +For all models, I report NDC@10 and queries per second. I do this for the student model and teacher model, to show how much performance you lose when switching between them. Detailed benchmark performance can be found in [the benchmarks folder.](./benchmarks/). The query timings were performed on the first 1000 queries of the msmarco dataset, and averaged over 7 runs. The benchmarks were run on an Apple M3 pro. + +### `gte-modernbert-base` + +| | Queries per second (CPU) | NDCG@10 | +|---------|--------------------------|---------| +| NIFE | 71400 (14ms/1k queries) | 59.2 | +| Teacher | 237 (4210ms/1k queries) | 66.34 | + +### `mxbai-embed-large-v1` + +| | Queries per second (CPU) | NDCG@10 | +|---------|--------------------------|---------| +| NIFE | 65789 (15ms/1k queries) | 59.2 | +| Teacher | 108 (9190ms/1k queries) | 65.6 | + +It is interesting that both NIFE models get the same performance, even with different teacher models. This could point towards a ceiling effect, where a certain percentage of queries can be answered correctly by static models, while others require contextualization. + ## How does it work? We use knowledge distillation from an initialized static model to the teacher we want to emulate. Some special things: @@ -147,156 +195,9 @@ NIFE can't do the following things: 1) Ignore words based on context: the query "What is the capital of France?" the word "France" will cause documents containing the term "France" to be retrieved. There is no way for the model to attenuate this vector and morph it into the answer ("Paris"). 2) Deal with negation: for the same reason as above; there is no interaction between tokens, so the similarity between "Cars that aren't red" and "Cars that are red" will be really high. -# Creating a NIFE model - -To create a NIFE model, you can run the scripts in `scripts`, or directly use the code from the repository. First, you should create a corpus of embeddings for your embedder. You can also use pre-computed collections of embeddings I created: - -* [mixedbread-ai/mxbai-embed-large-v1](https://huggingface.co/collections/stephantulkens/mxbai-large-v1-embedpress) -* [Alibaba-NLP/gte-modernbert-base](https://huggingface.co/collections/stephantulkens/gte-modernbert-embedpress) - -Broadly construed, training a NIFE model has 5 separate steps. - -### 1. Create a set of embeddings using the teacher - -Let's assume we want to create embeddings on [trivia QA](https://huggingface.co/mandarjoshi/trivia_qa), using `mxbai-embed-large-v1` as a teacher. - -```python -from datasets import load_dataset -from pynife.distillation.infer import generate_and_save_embeddings -from sentence_transformers import SentenceTransformer - -model_name = "mixedbread-ai/mxbai-embed-large-v1" -model = SentenceTransformer(model_name) - -dataset_name = "mandarjoshi/trivia_qa" -dataset = load_dataset(dataset_name, "rc", split="train") -dataset_iterator = (x['question'] for x in dataset) - -output_directory = "my-trivia-qa" - -generate_and_save_embeddings( - model=model, - records=dataset_iterator, - output_folder=output_directory, - limit_batches=None, - batch_size=8, - save_every=512, - max_length=512, - model_name=model_name, - dataset_name=dataset_name, - lowercase=False, - make_greedy=False, - ) - -``` - -This piece of code loads the model, the dataset and then starts inference. Inference takes a while, and will stream snippets to disk as .txt files and torch tensor files. After the whole dataset has been inferenced, the .txt and tensor files are converted into parquet files, and the .txt and torch tensor files are deleted. - -Your dataset will be ready and saved as parquet files in `output_directory`. If you want to upload these, please use the `HfAPI`, not `dataset.push_to_hub`, because we rely on some metadata embedded in the README to infer the base model later on. Note that the dataset iterator can be anything, and does not need to be a Hugging Face dataset. For example, it could also work with a stream from your database. - -For a simple inference script with a lot of pre-made datasets, see [the infer_datasets script](./scripts/infer_datasets.py). - -### 2. (optional) Expanding a tokenizer - -NIFE models work really well if you create a custom tokenizer for your domain. Empirically, it also works really well if you just expand the tokenizer of your teacher model with additional words. We call this _tokenizer expansion_. We have a pre-defined corpus to work on: - -```python -from transformers import AutoTokenizer - -from datasets import load_dataset -from pynife.tokenizer.expand_tokenizer import expand_tokenizer - - -dataset = load_dataset("stephantulkens/msmarco-vocab", split="train") -print(dataset.tolist()[:5]) -# [{'token': '.', 'frequency': 36174594, 'document_frequency': 8701009}, -# {'token': 'the', 'frequency': 28806701, 'document_frequency': 7712172}, -# {'token': ',', 'frequency': 25825435, 'document_frequency': 7411743}, -# {'token': 'of', 'frequency': 15196930, 'document_frequency': 6562023}, -# {'token': 'a', 'frequency': 13702107, 'document_frequency': 6064770}, - -tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") -# Function expects an iterator over dictionaries with "token" and "frequency" as keys. -new_tokenizer = expand_tokenizer(tokenizer, data, new_vocabulary_size=30000) -new_tokenizer.save_pretrained("my_tokenizer") - -``` - -This will do a couple of things: -1) It will remove all tokens from the original tokenizer that aren't present in your data. -2) It will then add the most frequent tokens until the size of the tokenizer == `new_vocabulary_size`. - -This works a lot better than training a tokenizer from scratch on equivalent data. For a runnable version, see [the expand_tokenizer script](./scripts/expand_tokenizer.py). - -To get frequency counts, you can use `count_tokens_in_dataset`, as follows: - -```python -from datasets import load_dataset, Dataset - -from pynife.tokenizer.count_vocabulary import count_tokens_in_dataset - -dataset = load_dataset("sentence-transformers/msmarco", "corpus", split="train", streaming=True) -dataset_iterator = (item["passage"] for item in dataset) -counts = count_tokens_in_dataset(dataset_iterator) - -# Save the counts as a dataset if you want. -dataset = Dataset.from_list(counts, split="train") -dataset.push_to_hub("my_hub") - -``` - -This dataset can be used directly to expand your tokenizer, above. For a runnable version, see [the create_vocabulary script](./scripts/create_vocabulary.py) - -### 3. Train - -Given a dataset and optionally a tokenizer, there's two steps to complete for a successful training. - -#### 3a Initialize a static model using your teacher - -Using *your teacher model*, initialize a static model. For example, when using [`mixedbread-ai/mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1): - -```python -from sentence_transformers import SentenceTransformer -from transformers import AutoTokenizer - -from pynife.initialization import initialize_from_model - -teacher = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") -# The tokenizer you trained in step 2. or an off-the-shelf tokenizer. -tokenizer = AutoTokenizer.from_pretrained("my_tokenizer") -model = initialize_from_model(teacher, tokenizer) - -``` - -#### 3b Actually train - -Now you can train, just like a regular sentence transformer. In my experiments, I found that using the cosine distance as a loss function was superior to using MSE, so I recommend using that, find it in `pynife.losses`. In addition, I also recommend using [Matryoshka Representation Learning](https://arxiv.org/abs/2205.13147). There's a bunch of helper functions in `pynife` to make training easier. In general, I recommend using hyperparameters like the following: - -* `batch_size`: 128 -* `learning rate`: 0.01 -* `scheduler`: "cosine_warmup_with_min_lr" -* `warmup_ratio`: 0.1 -* `weight_decay`: 0.01 -* `epochs`: 5 - -It can be tempting to move to very high batch sizes, but this has a very large detrimental effect on performance, even with higher learning rates. As a consequence, GPU usage during training is actually pretty low, because there's very little actual computation happening. For a complete runnable training loop, including model initialization, see [the training script](./scripts/experiment_distillation.py). - -```python -from pynife.losses import CosineLoss -from pynife.data import get_datasets - -# Fill with datasets you trained yourself. -datasets_you_made = [""] -train_dataset = get_datasets(datasets_you_made) - -# Model is initialized in step 3a. -loss = CosineLoss(model=model) - -# Train as usual. - -``` +## Inquiries -This will train a model and report the result to wandb. The `experiment_distillation` script is otherwise completely the same as a regular sentence transformers training loop, so there's very little actual code involved. +If you think NIFE could be interesting for your business let me know, I am open to consulting jobs regarding training models and fast inference. Just reach out to me [via e-mail](mailto:stephantul@gmail.com). ## License diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000..2451458 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,66 @@ +# Benchmarks + +This document contains detailed benchmarks for all nanoBEIR datasets, and some additional benchmarks with batch sizes. + +## Full nanoBEIR results + +Here are the full nanoBEIR results for each teacher/NIFE pair. + +### `NIFE-mxbai-embed-large-v1` + +| dataset | ndcg@10 teacher | ndcg@10 NIFE | +|:-------------------|------------------:|---------------:| +| NanoArguAna | 0.67 | 0.64 | +| NanoClimateFEVER | 0.40 | 0.32 | +| NanoDBPedia | 0.64 | 0.61 | +| NanoFEVER | 0.92 | 0.87 | +| NanoFiQA2018 | 0.56 | 0.44 | +| NanoHotpotQA | 0.87 | 0.75 | +| NanoMSMARCO | 0.66 | 0.62 | +| NanoNFCorpus | 0.38 | 0.34 | +| NanoNQ | 0.71 | 0.60 | +| NanoQuoraRetrieval | 0.96 | 0.90 | +| NanoSCIDOCS | 0.45 | 0.38 | +| NanoSciFact | 0.79 | 0.74 | +| NanoTouche2020 | 0.53 | 0.47 | + +The NIFE model performs worse on most datasets, but performs closely to its teacher on DBPedia and Arguana. Note that the NIFE model was trained on MSMARCO, so this represents a case of in-domain results: if you have many documents and can train on them, your NIFE model can approach the performance of your teacher model. + +### `NIFE-gte-modernbert-base` + +| dataset | ndcg@10 teacher | ndcg@10 NIFE | +|:-------------------|------------------:|---------------:| +| NanoArguAna | 0.77 | 0.61 | +| NanoClimateFEVER | 0.46 | 0.38 | +| NanoDBPedia | 0.60 | 0.61 | +| NanoFEVER | 0.94 | 0.81 | +| NanoFiQA2018 | 0.62 | 0.54 | +| NanoHotpotQA | 0.77 | 0.65 | +| NanoMSMARCO | 0.65 | 0.64 | +| NanoNFCorpus | 0.35 | 0.35 | +| NanoNQ | 0.72 | 0.63 | +| NanoQuoraRetrieval | 0.97 | 0.90 | +| NanoSCIDOCS | 0.47 | 0.37 | +| NanoSciFact | 0.82 | 0.75 | +| NanoTouche2020 | 0.48 | 0.46 | + +As you can see the NIFE model performs worse on most datasets, but outperforms the base model on DBPedia and is pretty close on NanoMSMARCO. Note that the NIFE model was trained on MSMARCO, so this represents a case of in-domain results: if you have many documents and can train on them, your NIFE model can approach the performance of your teacher model. + +## Effect of batch size on speed + +One interesting issue is that most query embedders are not run on very large batch sizes. In real-life workloads, embedders are often run in microservices that ingest a single query at a time. Therefore, reporting QPS at high batch sizes is actually not a realistic estimate of performance. + +For this experiment, I use `NIFE-mxbai-embed-large-v1`, all timings are done on CPU (Macbook Pro M3) + +| Batch size | NIFE QPS | Teacher QPS | x Speedup | +|----:|-------:|----------:|------------:| +| 1 | 12925 | 13 | 1007 | +| 2 | 14090 | 26 | 552 | +| 4 | 19670 | 45 | 435 | +| 8 | 31326 | 64 | 486 | +| 16 | 46346 | 82 | 567 | +| 32 | 62077 | 99 | 628 | +| 64 | 75661 | 107 | 706 | +| 128 | 93948 | 101 | 931 | + +As you can see both NIFE and the teacher benefit a lot from batching, although NIFE is already very fast to begin with. This shows that, on CPU, there is a very large gain on QPS by switching from the teacher to NIFE. diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..aa55d36 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,154 @@ +# Documentation + +This section currently just includes documentation on how to create your own NIFE model. + +## Creating a NIFE model + +To create a NIFE model, you can run the scripts in `scripts`, or directly use the code from the repository. First, you should create a corpus of embeddings for your embedder. You can also use pre-computed collections of embeddings I created: + +* [mixedbread-ai/mxbai-embed-large-v1](https://huggingface.co/collections/stephantulkens/mxbai-large-v1-embedpress) +* [Alibaba-NLP/gte-modernbert-base](https://huggingface.co/collections/stephantulkens/gte-modernbert-embedpress) + +Broadly construed, training a NIFE model has 5 separate steps. + +### 1. Create a set of embeddings using the teacher + +Let's assume we want to create embeddings on [trivia QA](https://huggingface.co/mandarjoshi/trivia_qa), using `mxbai-embed-large-v1` as a teacher. + +```python +from datasets import load_dataset +from pynife.distillation.infer import generate_and_save_embeddings +from sentence_transformers import SentenceTransformer + +model_name = "mixedbread-ai/mxbai-embed-large-v1" +model = SentenceTransformer(model_name) + +dataset_name = "mandarjoshi/trivia_qa" +dataset = load_dataset(dataset_name, "rc", split="train") +dataset_iterator = (x['question'] for x in dataset) + +output_directory = "my-trivia-qa" + +generate_and_save_embeddings( + model=model, + records=dataset_iterator, + output_folder=output_directory, + limit_batches=None, + batch_size=8, + save_every=512, + max_length=512, + model_name=model_name, + dataset_name=dataset_name, + lowercase=False, + make_greedy=False, + ) + +``` + +This piece of code loads the model, the dataset and then starts inference. Inference takes a while, and will stream snippets to disk as .txt files and torch tensor files. After the whole dataset has been inferenced, the .txt and tensor files are converted into parquet files, and the .txt and torch tensor files are deleted. + +Your dataset will be ready and saved as parquet files in `output_directory`. If you want to upload these, please use the `HfAPI`, not `dataset.push_to_hub`, because we rely on some metadata embedded in the README to infer the base model later on. Note that the dataset iterator can be anything, and does not need to be a Hugging Face dataset. For example, it could also work with a stream from your database. + +For a simple inference script with a lot of pre-made datasets, see [the infer_datasets script](./scripts/infer_datasets.py). + +### 2. (optional) Expanding a tokenizer + +NIFE models work really well if you create a custom tokenizer for your domain. Empirically, it also works really well if you just expand the tokenizer of your teacher model with additional words. We call this _tokenizer expansion_. We have a pre-defined corpus to work on: + +```python +from transformers import AutoTokenizer + +from datasets import load_dataset +from pynife.tokenizer.expand_tokenizer import expand_tokenizer + + +dataset = load_dataset("stephantulkens/msmarco-vocab", split="train") +print(dataset.tolist()[:5]) +# [{'token': '.', 'frequency': 36174594, 'document_frequency': 8701009}, +# {'token': 'the', 'frequency': 28806701, 'document_frequency': 7712172}, +# {'token': ',', 'frequency': 25825435, 'document_frequency': 7411743}, +# {'token': 'of', 'frequency': 15196930, 'document_frequency': 6562023}, +# {'token': 'a', 'frequency': 13702107, 'document_frequency': 6064770}, + +tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") +# Function expects an iterator over dictionaries with "token" and "frequency" as keys. +new_tokenizer = expand_tokenizer(tokenizer, data, new_vocabulary_size=30000) +new_tokenizer.save_pretrained("my_tokenizer") + +``` + +This will do a couple of things: +1) It will remove all tokens from the original tokenizer that aren't present in your data. +2) It will then add the most frequent tokens until the size of the tokenizer == `new_vocabulary_size`. + +This works a lot better than training a tokenizer from scratch on equivalent data. For a runnable version, see [the expand_tokenizer script](./scripts/expand_tokenizer.py). + +To get frequency counts, you can use `count_tokens_in_dataset`, as follows: + +```python +from datasets import load_dataset, Dataset + +from pynife.tokenizer.count_vocabulary import count_tokens_in_dataset + +dataset = load_dataset("sentence-transformers/msmarco", "corpus", split="train", streaming=True) +dataset_iterator = (item["passage"] for item in dataset) +counts = count_tokens_in_dataset(dataset_iterator) + +# Save the counts as a dataset if you want. +dataset = Dataset.from_list(counts, split="train") +dataset.push_to_hub("my_hub") + +``` + +This dataset can be used directly to expand your tokenizer, above. For a runnable version, see [the create_vocabulary script](./scripts/create_vocabulary.py) + +### 3. Train + +Given a dataset and optionally a tokenizer, there's two steps to complete for a successful training. + +#### 3a Initialize a static model using your teacher + +Using *your teacher model*, initialize a static model. For example, when using [`mixedbread-ai/mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1): + +```python +from sentence_transformers import SentenceTransformer +from transformers import AutoTokenizer + +from pynife.initialization import initialize_from_model + +teacher = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") +# The tokenizer you trained in step 2. or an off-the-shelf tokenizer. +tokenizer = AutoTokenizer.from_pretrained("my_tokenizer") +model = initialize_from_model(teacher, tokenizer) + +``` + +#### 3b Actually train + +Now you can train, just like a regular sentence transformer. In my experiments, I found that using the cosine distance as a loss function was superior to using MSE, so I recommend using that, find it in `pynife.losses`. In addition, I also recommend using [Matryoshka Representation Learning](https://arxiv.org/abs/2205.13147). There's a bunch of helper functions in `pynife` to make training easier. In general, I recommend using hyperparameters like the following: + +* `batch_size`: 128 +* `learning rate`: 0.01 +* `scheduler`: "cosine_warmup_with_min_lr" +* `warmup_ratio`: 0.1 +* `weight_decay`: 0.01 +* `epochs`: 5 + +It can be tempting to move to very high batch sizes, but this has a very large detrimental effect on performance, even with higher learning rates. As a consequence, GPU usage during training is actually pretty low, because there's very little actual computation happening. For a complete runnable training loop, including model initialization, see [the training script](./scripts/experiment_distillation.py). + +```python +from pynife.losses import CosineLoss +from pynife.data import get_datasets + +# Fill with datasets you trained yourself. +datasets_you_made = [""] +train_dataset = get_datasets(datasets_you_made) + +# Model is initialized in step 3a. +loss = CosineLoss(model=model) + +# Train as usual. + +``` + +This will train a model and report the result to wandb. The `experiment_distillation` script is otherwise completely the same as a regular sentence transformers training loop, so there's very little actual code involved.