PyLate model based on answerdotai/ModernBERT-base

This is a PyLate model finetuned from answerdotai/ModernBERT-base on the msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1 dataset. It maps sentences & paragraphs to sequences of 128-dimensional dense vectors and can be used for semantic textual similarity using the MaxSim operator.

Model Details

Model Description

Model Type: PyLate model
Base model: answerdotai/ModernBERT-base
Document Length: 180 tokens
Query Length: 32 tokens
Output Dimensionality: 128 tokens
Similarity Function: MaxSim
Training Dataset:
- msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1
Language: en

Model Sources

Documentation: PyLate Documentation
Repository: PyLate on GitHub
Hugging Face: PyLate models on Hugging Face

Full Model Architecture

ColBERT(
  (0): Transformer({'max_seq_length': 8192, 'do_lower_case': False, 'architecture': 'ModernBertModel'})
  (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'use_residual': False})
)

Usage

First install the PyLate library:

pip install -U pylate

Retrieval

Use this model with PyLate to index and retrieve documents. The index uses FastPLAID for efficient similarity search.

Indexing documents

Load the ColBERT model and initialize the PLAID index, then encode and index your documents:

from pylate import indexes, models, retrieve

# Step 1: Load the ColBERT model
model = models.ColBERT(
    model_name_or_path="PrasannSinghal/ModernBERT-base-pylate-pairwise-8e-05-qv1-dv1-embsize128",
)

# Step 2: Initialize the PLAID index
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # This overwrites the existing index if any
)

# Step 3: Encode the documents
documents_ids = ["1", "2", "3"]
documents = ["document 1 text", "document 2 text", "document 3 text"]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,  # Ensure that it is set to False to indicate that these are documents, not queries
    show_progress_bar=True,
)

# Step 4: Add document embeddings to the index by providing embeddings and corresponding ids
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

Note that you do not have to recreate the index and encode the documents every time. Once you have created an index and added the documents, you can re-use the index later by loading it:

# To load an index, simply instantiate it with the correct folder/name and without overriding it
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
)

Retrieving top-k documents for queries

Once the documents are indexed, you can retrieve the top-k most relevant documents for a given set of queries. To do so, initialize the ColBERT retriever with the index you want to search in, encode the queries and then retrieve the top-k documents to get the top matches ids and relevance scores:

# Step 1: Initialize the ColBERT retriever
retriever = retrieve.ColBERT(index=index)

# Step 2: Encode the queries
queries_embeddings = model.encode(
    ["query for document 3", "query for document 1"],
    batch_size=32,
    is_query=True,  #  # Ensure that it is set to False to indicate that these are queries
    show_progress_bar=True,
)

# Step 3: Retrieve top-k documents
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=10,  # Retrieve the top 10 matches for each query
)

Reranking

If you only want to use the ColBERT model to perform reranking on top of your first-stage retrieval pipeline without building an index, you can simply use rank function and pass the queries and documents to rerank:

from pylate import rank, models

queries = [
    "query A",
    "query B",
]

documents = [
    ["document A", "document B"],
    ["document 1", "document C", "document B"],
]

documents_ids = [
    [1, 2],
    [1, 3, 2],
]

model = models.ColBERT(
    model_name_or_path="PrasannSinghal/ModernBERT-base-pylate-pairwise-8e-05-qv1-dv1-embsize128",
)

queries_embeddings = model.encode(
    queries,
    is_query=True,
)

documents_embeddings = model.encode(
    documents,
    is_query=False,
)

reranked_documents = rank.rerank(
    documents_ids=documents_ids,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings,
)

Evaluation

Metrics

Col BERTTriplet

Dataset: msmarco-co-condenser-dev
Evaluated with pylate.evaluation.colbert_triplet.ColBERTTripletEvaluator

Metric	Value
accuracy	0.864

Training Details

Training Dataset

msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1

Dataset: msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1 at 84ed2d3
Size: 1,250,000 training samples
Columns: query, positive, and negative

Approximate statistics based on the first 1000 samples:

	query	positive	negative
type	string	string	string
details	min: 7 tokens mean: 12.26 tokens max: 37 tokens	min: 20 tokens mean: 82.14 tokens max: 225 tokens	min: 27 tokens mean: 83.09 tokens max: 439 tokens

Samples:

query	positive	negative
`what is the meaning of menu planning`	`Menu planning is the selection of a menu for an event. Such as picking out the dinner for your wedding or even a meal at a Birthday Party. Menu planning is when you are preparing a calendar of meals and you have to sit down and decide what meat and veggies you want to serve on each certain day.`	`Menu Costs. In economics, a menu cost is the cost to a firm resulting from changing its prices. The name stems from the cost of restaurants literally printing new menus, but economists use it to refer to the costs of changing nominal prices in general.`
`how old is brett butler`	`Brett Butler is 59 years old. To be more precise (and nerdy), the current age as of right now is 21564 days or (even more geeky) 517536 hours. That's a lot of hours!`	`Passed in: St. John's, Newfoundland and Labrador, Canada. Passed on: 16/07/2016. Published in the St. John's Telegram. Passed away suddenly at the Health Sciences Centre surrounded by his loving family, on July 16, 2016 Robert (Bobby) Joseph Butler, age 52 years. Predeceased by his special aunt Geri Murrin and uncle Mike Mchugh; grandparents Joe and Margaret Murrin and Jack and Theresa Butler.`
`when was the last navajo treaty sign?`	`In Executive Session, Senate of the United States, July 25, 1868. Resolved, (two-thirds of the senators present concurring,) That the Senate advise and consent to the ratification of the treaty between the United States and the Navajo Indians, concluded at Fort Sumner, New Mexico, on the first day of June, 1868.`	`Share Treaty of Greenville. The Treaty of Greenville was signed August 3, 1795, between the United States, represented by Gen. Anthony Wayne, and chiefs of the Indian tribes located in the Northwest Territory, including the Wyandots, Delawares, Shawnees, Ottawas, Miamis, and others.`

Loss: pylate.losses.cached_contrastive.CachedContrastive

Evaluation Dataset

msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1

Dataset: msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1 at 84ed2d3
Size: 1,000 evaluation samples
Columns: query, positive, and negative

Approximate statistics based on the first 1000 samples:

	query	positive	negative
type	string	string	string
details	min: 7 tokens mean: 12.2 tokens max: 30 tokens	min: 24 tokens mean: 83.44 tokens max: 244 tokens	min: 26 tokens mean: 83.38 tokens max: 242 tokens

Samples:

query	positive	negative
`what county is holly springs nc in`	`Holly Springs, North Carolina. Holly Springs is a town in Wake County, North Carolina, United States. As of the 2010 census, the town population was 24,661, over 2Â½ times its population in 2000. Contents.`	`The Mt. Holly Springs Park & Resort. One of the numerous trolley routes that carried people around the county at the turn of the century was the Carlisle & Mt. Holly Railway Company. The âHolly Trolleyâ as it came to be known was put into service by Patricio Russo and made its first run on May 14, 1901.`
`how long does nyquil stay in your system`	In order to understand exactly how long Nyquil lasts, it is absolutely vital to learn about the various ingredients in the drug. One of the ingredients found in Nyquil is Doxylamine, which is an antihistamine. This specific medication has a biological half-life or 6 to 12 hours. With this in mind, it is possible for the drug to remain in the system for a period of 12 to 24 hours. It should be known that the specifics will depend on a wide variety of different factors, including your age and metabolism.	I confirmed that NyQuil is about 10% alcohol, a higher content than most domestic beers. When I asked about the relatively high proof, I was told that the alcohol dilutes the active ingredients. The alcohol free version is there for customers with addiction issues.. also found that in that version there is twice the amount of DXM. When I asked if I could speak to a chemist or scientist, I was told they didn't have anyone who fit that description there. Itâs been eight years since I kicked NyQuil. I've been sober from alcohol for four years.
`what are mineral water`	1 Mineral water â water from a mineral spring that contains various minerals, such as salts and sulfur compounds. 2 It comes from a source tapped at one or more bore holes or spring, and originates from a geologically and physically protected underground water source. Mineral water â water from a mineral spring that contains various minerals, such as salts and sulfur compounds. 2 It comes from a source tapped at one or more bore holes or spring, and originates from a geologically and physically protected underground water source.	`Minerals for Your Body. Drinking mineral water is beneficial to health and well-being. But it is not only the amount of water you drink that is important-what the water contains is even more essential.inerals for Your Body. Drinking mineral water is beneficial to health and well-being. But it is not only the amount of water you drink that is important-what the water contains is even more essential.`

Loss: pylate.losses.cached_contrastive.CachedContrastive

Training Hyperparameters

Non-Default Hyperparameters

per_device_train_batch_size: 512
per_device_eval_batch_size: 512
learning_rate: 8e-05
num_train_epochs: 1
warmup_ratio: 0.05
bf16: True
dataloader_num_workers: 8
dataloader_pin_memory: False
batch_sampler: no_duplicates

All Hyperparameters

Click to expand

overwrite_output_dir: False
do_predict: False
eval_strategy: no
prediction_loss_only: True
per_device_train_batch_size: 512
per_device_eval_batch_size: 512
per_gpu_train_batch_size: None
per_gpu_eval_batch_size: None
gradient_accumulation_steps: 1
eval_accumulation_steps: None
torch_empty_cache_steps: None
learning_rate: 8e-05
weight_decay: 0.0
adam_beta1: 0.9
adam_beta2: 0.999
adam_epsilon: 1e-08
max_grad_norm: 1.0
num_train_epochs: 1
max_steps: -1
lr_scheduler_type: linear
lr_scheduler_kwargs: {}
warmup_ratio: 0.05
warmup_steps: 0
log_level: passive
log_level_replica: warning
log_on_each_node: True
logging_nan_inf_filter: True
save_safetensors: True
save_on_each_node: False
save_only_model: False
restore_callback_states_from_checkpoint: False
no_cuda: False
use_cpu: False
use_mps_device: False
seed: 42
data_seed: None
jit_mode_eval: False
use_ipex: False
bf16: True
fp16: False
fp16_opt_level: O1
half_precision_backend: auto
bf16_full_eval: False
fp16_full_eval: False
tf32: None
local_rank: 1
ddp_backend: None
tpu_num_cores: None
tpu_metrics_debug: False
debug: []
dataloader_drop_last: True
dataloader_num_workers: 8
dataloader_prefetch_factor: None
past_index: -1
disable_tqdm: False
remove_unused_columns: True
label_names: None
load_best_model_at_end: False
ignore_data_skip: False
fsdp: []
fsdp_min_num_params: 0
fsdp_config: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}
fsdp_transformer_layer_cls_to_wrap: None
accelerator_config: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}
parallelism_config: None
deepspeed: None
label_smoothing_factor: 0.0
optim: adamw_torch_fused
optim_args: None
adafactor: False
group_by_length: False
length_column_name: length
ddp_find_unused_parameters: None
ddp_bucket_cap_mb: None
ddp_broadcast_buffers: False
dataloader_pin_memory: False
dataloader_persistent_workers: False
skip_memory_metrics: True
use_legacy_prediction_loop: False
push_to_hub: False
resume_from_checkpoint: None
hub_model_id: None
hub_strategy: every_save
hub_private_repo: None
hub_always_push: False
hub_revision: None
gradient_checkpointing: False
gradient_checkpointing_kwargs: None
include_inputs_for_metrics: False
include_for_metrics: []
eval_do_concat_batches: True
fp16_backend: auto
push_to_hub_model_id: None
push_to_hub_organization: None
mp_parameters:
auto_find_batch_size: False
full_determinism: False
torchdynamo: None
ray_scope: last
ddp_timeout: 1800
torch_compile: False
torch_compile_backend: None
torch_compile_mode: None
include_tokens_per_second: False
include_num_input_tokens_seen: False
neftune_noise_alpha: None
optim_target_modules: None
batch_eval_metrics: False
eval_on_start: False
use_liger_kernel: False
liger_kernel_config: None
eval_use_gather_object: False
average_tokens_across_devices: True
prompts: None
batch_sampler: no_duplicates
multi_dataset_batch_sampler: proportional
router_mapping: {}
learning_rate_mapping: {}

Training Logs

Click to expand

Epoch	Step	Training Loss	accuracy
0	0	-	0.5070
0.0049	2	7.8753	-
0.0099	4	7.7978	-
0.0148	6	7.8206	-
0.0197	8	7.5634	-
0.0246	10	7.3687	-
0.0296	12	7.2176	-
0.0345	14	7.0843	-
0.0394	16	7.0076	-
0.0443	18	6.9584	-
0.0493	20	6.9296	-
0.0542	22	6.9121	-
0.0591	24	6.9004	-
0.0640	26	6.8907	-
0.0690	28	6.8603	-
0.0739	30	6.815	-
0.0788	32	6.7927	-
0.0837	34	6.7439	-
0.0887	36	6.6842	-
0.0936	38	6.6257	-
0.0985	40	6.522	-
0.1034	42	6.4258	-
0.1084	44	6.2891	-
0.1133	46	6.1493	-
0.1182	48	5.9112	-
0.1232	50	5.5444	-
0.1281	52	5.0781	-
0.1330	54	4.5032	-
0.1379	56	4.031	-
0.1429	58	3.5449	-
0.1478	60	3.0924	-
0.1527	62	2.7258	-
0.1576	64	2.4606	-
0.1626	66	2.2058	-
0.1675	68	2.0276	-
0.1724	70	1.8286	-
0.1773	72	1.7153	-
0.1823	74	1.6934	-
0.1872	76	1.5035	-
0.1921	78	1.4719	-
0.1970	80	1.401	-
0.2020	82	1.3597	-
0.2069	84	1.3274	-
0.2118	86	1.299	-
0.2167	88	1.2498	-
0.2217	90	1.198	-
0.2266	92	1.1907	-
0.2315	94	1.154	-
0.2365	96	1.1252	-
0.2414	98	1.1108	-
0.2463	100	1.1037	-
0.2512	102	1.0542	-
0.2562	104	1.0743	-
0.2611	106	1.0335	-
0.2660	108	1.015	-
0.2709	110	0.9684	-
0.2759	112	0.9739	-
0.2808	114	0.9918	-
0.2857	116	0.9793	-
0.2906	118	0.9383	-
0.2956	120	0.935	-
0.3005	122	0.9261	-
0.3054	124	0.9262	-
0.3103	126	0.914	-
0.3153	128	0.8921	-
0.3202	130	0.8767	-
0.3251	132	0.9278	-
0.3300	134	0.8952	-
0.3350	136	0.8258	-
0.3399	138	0.8793	-
0.3448	140	0.8431	-
0.3498	142	0.8199	-
0.3547	144	0.8188	-
0.3596	146	0.828	-
0.3645	148	0.8203	-
0.3695	150	0.8099	-
0.3744	152	0.8151	-
0.3793	154	0.814	-
0.3842	156	0.7808	-
0.3892	158	0.8124	-
0.3941	160	0.8023	-
0.3990	162	0.7805	-
0.4039	164	0.781	-
0.4089	166	0.7624	-
0.4138	168	0.7646	-
0.4187	170	0.7642	-
0.4236	172	0.7593	-
0.4286	174	0.7598	-
0.4335	176	0.7652	-
0.4384	178	0.7416	-
0.4433	180	0.7028	-
0.4483	182	0.7317	-
0.4532	184	0.7312	-
0.4581	186	0.7273	-
0.4631	188	0.7114	-
0.4680	190	0.727	-
0.4729	192	0.7175	-
0.4778	194	0.7288	-
0.4828	196	0.7147	-
0.4877	198	0.6986	-
0.4926	200	0.7023	-
0.4975	202	0.6991	-
0.5025	204	0.7056	-
0.5074	206	0.7157	-
0.5123	208	0.6935	-
0.5172	210	0.6991	-
0.5222	212	0.6963	-
0.5271	214	0.666	-
0.5320	216	0.6752	-
0.5369	218	0.6973	-
0.5419	220	0.7014	-
0.5468	222	0.6436	-
0.5517	224	0.6636	-
0.5567	226	0.6811	-
0.5616	228	0.6643	-
0.5665	230	0.6835	-
0.5714	232	0.6714	-
0.5764	234	0.6666	-
0.5813	236	0.6589	-
0.5862	238	0.6616	-
0.5911	240	0.6678	-
0.5961	242	0.6566	-
0.6010	244	0.6513	-
0.6059	246	0.6704	-
0.6108	248	0.6711	-
0.6158	250	0.6348	-
0.6207	252	0.6711	-
0.6256	254	0.6318	-
0.6305	256	0.6455	-
0.6355	258	0.6285	-
0.6404	260	0.6456	-
0.6453	262	0.641	-
0.6502	264	0.6506	-
0.6552	266	0.6362	-
0.6601	268	0.6492	-
0.6650	270	0.6409	-
0.6700	272	0.6264	-
0.6749	274	0.6223	-
0.6798	276	0.6318	-
0.6847	278	0.6264	-
0.6897	280	0.6306	-
0.6946	282	0.664	-
0.6995	284	0.6061	-
0.7044	286	0.6488	-
0.7094	288	0.6335	-
0.7143	290	0.6695	-
0.7192	292	0.6251	-
0.7241	294	0.6396	-
0.7291	296	0.6089	-
0.7340	298	0.6319	-
0.7389	300	0.6576	-
0.7438	302	0.6164	-
0.7488	304	0.6276	-
0.7537	306	0.616	-
0.7586	308	0.6436	-
0.7635	310	0.5965	-
0.7685	312	0.6225	-
0.7734	314	0.6352	-
0.7783	316	0.6247	-
0.7833	318	0.6374	-
0.7882	320	0.6091	-
0.7931	322	0.624	-
0.7980	324	0.6293	-
0.8030	326	0.6061	-
0.8079	328	0.6425	-
0.8128	330	0.6372	-
0.8177	332	0.6048	-
0.8227	334	0.6152	-
0.8276	336	0.6131	-
0.8325	338	0.6321	-
0.8374	340	0.6318	-
0.8424	342	0.6215	-
0.8473	344	0.594	-
0.8522	346	0.6327	-
0.8571	348	0.5966	-
0.8621	350	0.5936	-
0.8670	352	0.63	-
0.8719	354	0.6102	-
0.8768	356	0.6144	-
0.8818	358	0.6097	-
0.8867	360	0.6097	-
0.8916	362	0.602	-
0.8966	364	0.6347	-
0.9015	366	0.6076	-
0.9064	368	0.6195	-
0.9113	370	0.5991	-
0.9163	372	0.6301	-
0.9212	374	0.6263	-
0.9261	376	0.613	-
0.9310	378	0.6385	-
0.9360	380	0.6092	-
0.9409	382	0.6178	-
0.9458	384	0.6049	-
0.9507	386	0.6099	-
0.9557	388	0.6317	-
0.9606	390	0.6207	-
0.9655	392	0.6045	-
0.9704	394	0.5905	-
0.9754	396	0.6199	-
0.9803	398	0.6181	-
0.9852	400	0.6077	-
0.9901	402	0.6325	-
0.9951	404	0.6161	-
1.0	406	0.6209	-
0	0	-	0.8640

Framework Versions

Python: 3.10.0
Sentence Transformers: 5.2.2
PyLate: 1.3.4
Transformers: 4.56.2
PyTorch: 2.9.1+cu128
Accelerate: 1.12.0
Datasets: 4.5.0
Tokenizers: 0.22.2

Citation

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084"
}

PyLate

@misc{PyLate,
title={PyLate: Flexible Training and Retrieval for Late Interaction Models},
author={Chaffin, Antoine and Sourty, Raphaël},
url={https://github.com/lightonai/pylate},
year={2024}
}

CachedContrastive

@misc{gao2021scaling,
    title={Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup},
    author={Luyu Gao and Yunyi Zhang and Jiawei Han and Jamie Callan},
    year={2021},
    eprint={2101.06983},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}