You can deploy this model and start earning money today!

Views

No views yet

Fine-tuned ColBERT model for semantic caching

Name: langcache-colbert-v1
Availability: InStock

This is a PyLate model finetuned from colbert-ir/colbertv2.0 on the LangCache Sentence Pairs (subsets=['all'], train+val=True) 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: colbert-ir/colbertv2.0
Document Length: 128 tokens
Query Length: 128 tokens
Output Dimensionality: 128 tokens
Similarity Function: MaxSim
Training Dataset:
- LangCache Sentence Pairs (subsets=['all'], train+val=True)
Language: en
License: apache-2.0

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': 127, 'do_lower_case': False, 'architecture': 'BertModel'})
  (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:

python

from pylate import indexes, models, retrieve

# Step 1: Load the ColBERT model
model = models.ColBERT(
    model_name_or_path="aditeyabaral/langcache-colbert-v1",
)

# 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:

python

# 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:

python

# 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:

python

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="aditeyabaral/langcache-colbert-v1",
)

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: test_triplet
Evaluated with pylate.evaluation.colbert_triplet.ColBERTTripletEvaluator

Metric	Value
accuracy	0.9847

Training Details

Training Dataset

LangCache Sentence Pairs (subsets=['all'], train+val=True)

Dataset: LangCache Sentence Pairs (subsets=['all'], train+val=True)
Size: 1,452,533 training samples
Columns: anchor, positive, and negative_1
Approximate statistics based on the first 1000 samples:
anchor positive negative_1
type string string string
details
min: 9 tokens
mean: 29.49 tokens
max: 73 tokens
min: 8 tokens
mean: 29.18 tokens
max: 58 tokens
min: 4 tokens
mean: 23.79 tokens
max: 52 tokens

Samples:

anchor	positive	negative_1
`Any Canadian teachers (B.Ed. holders) teaching in U.S. schools?`	`Any Canadian teachers (B.Ed. holders) teaching in U.S. schools?`	`Are there many Canadians living and working illegally in the United States?`
`Are there any underlying psychological tricks/tactics that are used when designing the lines for rides at amusement parks?`	`Are there any underlying psychological tricks/tactics that are used when designing the lines for rides at amusement parks?`	`Is there any tricks for straight lines mcqs?`
`Can I pay with a debit card on PayPal?`	`Can I pay with a debit card on PayPal?`	`Can you transfer PayPal funds onto a debit card/credit card?`

Loss: pylate.losses.contrastive.Contrastive

Evaluation Dataset

LangCache Sentence Pairs (split=test)

Dataset: LangCache Sentence Pairs (split=test)
Size: 110,066 evaluation samples
Columns: anchor, positive, and negative_1
Approximate statistics based on the first 1000 samples:
anchor positive negative_1
type string string string
details
min: 5 tokens
mean: 28.57 tokens
max: 121 tokens
min: 5 tokens
mean: 28.01 tokens
max: 121 tokens
min: 7 tokens
mean: 20.73 tokens
max: 65 tokens

Samples:

anchor	positive	negative_1
`What high potential jobs are there other than computer science?`	`What high potential jobs are there other than computer science?`	`Why IT or Computer Science jobs are being over rated than other Engineering jobs?`
`Would India ever be able to develop a missile system like S300 or S400 missile?`	`Would India ever be able to develop a missile system like S300 or S400 missile?`	`Should India buy the Russian S400 air defence missile system?`
`water from the faucet is being drunk by a yellow dog`	`A yellow dog is drinking water from the faucet`	`Do you get more homework in 9th grade than 8th?`

Loss: pylate.losses.contrastive.Contrastive

Training Hyperparameters

Non-Default Hyperparameters

per_device_train_batch_size: 48
num_train_epochs: 5
learning_rate: 0.0002
warmup_steps: 0.1
optim: adamw_torch
weight_decay: 0.001
eval_strategy: steps
per_device_eval_batch_size: 48
eval_on_start: True
push_to_hub: True
hub_model_id: aditeyabaral/langcache-colbert-v1
load_best_model_at_end: True
ddp_find_unused_parameters: True
batch_sampler: no_duplicates

All Hyperparameters

Click to expand

per_device_train_batch_size: 48
num_train_epochs: 5
max_steps: -1
learning_rate: 0.0002
lr_scheduler_type: linear
lr_scheduler_kwargs: None
warmup_steps: 0.1
optim: adamw_torch
optim_args: None
weight_decay: 0.001
adam_beta1: 0.9
adam_beta2: 0.999
adam_epsilon: 1e-08
optim_target_modules: None
gradient_accumulation_steps: 1
average_tokens_across_devices: True
max_grad_norm: 1.0
label_smoothing_factor: 0.0
bf16: False
fp16: False
bf16_full_eval: False
fp16_full_eval: False
tf32: None
gradient_checkpointing: False
gradient_checkpointing_kwargs: None
torch_compile: False
torch_compile_backend: None
torch_compile_mode: None
use_liger_kernel: False
liger_kernel_config: None
use_cache: False
neftune_noise_alpha: None
torch_empty_cache_steps: None
auto_find_batch_size: False
log_on_each_node: True
logging_nan_inf_filter: True
include_num_input_tokens_seen: no
log_level: passive
log_level_replica: warning
disable_tqdm: False
project: huggingface
trackio_space_id: trackio
eval_strategy: steps
per_device_eval_batch_size: 48
prediction_loss_only: True
eval_on_start: True
eval_do_concat_batches: True
eval_use_gather_object: False
eval_accumulation_steps: None
include_for_metrics: []
batch_eval_metrics: False
save_only_model: False
save_on_each_node: False
enable_jit_checkpoint: False
push_to_hub: True
hub_private_repo: None
hub_model_id: aditeyabaral/langcache-colbert-v1
hub_strategy: every_save
hub_always_push: False
hub_revision: None
load_best_model_at_end: True
ignore_data_skip: False
restore_callback_states_from_checkpoint: False
full_determinism: False
seed: 42
data_seed: None
use_cpu: False
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
dataloader_drop_last: True
dataloader_num_workers: 0
dataloader_pin_memory: True
dataloader_persistent_workers: False
dataloader_prefetch_factor: None
remove_unused_columns: True
label_names: None
train_sampling_strategy: random
length_column_name: length
ddp_find_unused_parameters: True
ddp_bucket_cap_mb: None
ddp_broadcast_buffers: False
ddp_backend: None
ddp_timeout: 1800
fsdp: []
fsdp_config: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}
deepspeed: None
debug: []
skip_memory_metrics: True
do_predict: False
resume_from_checkpoint: None
warmup_ratio: None
local_rank: -1
prompts: None
batch_sampler: no_duplicates
multi_dataset_batch_sampler: proportional
router_mapping: {}
learning_rate_mapping: {}

Training Logs

Epoch	Step	Training Loss	Validation Loss	accuracy
0	0	-	527.3283	0.8206
0.0661	1000	73.6149	-	-
0.1322	2000	5.4919	-	-
0.1983	3000	0.3036	-	-
0.2644	4000	0.2963	-	-
0.3305	5000	0.3388	-	-
0.3966	6000	0.2512	-	-
0.4627	7000	0.2497	-	-
0.5288	8000	0.2427	-	-
0.5948	9000	0.2585	-	-
0.6609	10000	0.5272	-	-
0.7270	11000	0.2143	-	-
0.7931	12000	0.2065	-	-
0.8592	13000	0.4024	-	-
0.9253	14000	0.6400	-	-
0.9914	15000	0.8233	-	-
1.0575	16000	0.7687	-	-
1.1236	17000	0.7550	-	-
1.1897	18000	0.6507	-	-
1.2558	19000	0.6809	-	-
1.3219	20000	0.6523	-	-
1.3880	21000	0.5745	-	-
1.4541	22000	0.5485	-	-
1.5202	23000	0.5092	-	-
1.5863	24000	0.4815	-	-
1.6523	25000	0.4785	-	-
1.7184	26000	0.4901	-	-
1.7845	27000	0.4581	-	-
1.8506	28000	0.5224	-	-
1.9167	29000	0.4892	-	-
1.9828	30000	0.4884	-	-
2.0489	31000	0.4530	-	-
2.1150	32000	0.4356	-	-
2.1811	33000	0.4555	-	-
2.2472	34000	0.4360	-	-
2.3133	35000	0.4478	-	-
2.3794	36000	0.4297	-	-
2.4455	37000	0.3896	-	-
2.5116	38000	0.3594	-	-
2.5777	39000	0.3581	-	-
2.6438	40000	0.3270	-	-
2.7098	41000	0.3995	-	-
2.7759	42000	0.3665	-	-
2.8420	43000	0.4018	-	-
2.9081	44000	0.4260	-	-
2.9742	45000	0.3957	-	-
3.0403	46000	0.3659	-	-
3.1064	47000	0.3826	-	-
3.1725	48000	0.3603	-	-
3.2386	49000	0.3646	-	-
3.3047	50000	0.4069	-	-
0	0	-	-	0.9847
3.3047	50000	-	2.1447	-
3.3708	51000	0.3493	-	-
3.4369	52000	0.3207	-	-
3.5030	53000	0.3311	-	-
3.5691	54000	0.3208	-	-
3.6352	55000	0.2760	-	-
3.7013	56000	0.3244	-	-
3.7673	57000	0.2789	-	-
3.8334	58000	0.3038	-	-
3.8995	59000	0.3958	-	-
3.9656	60000	0.3338	-	-
4.0317	61000	0.3445	-	-
4.0978	62000	0.3291	-	-
4.1639	63000	0.3225	-	-
4.2300	64000	0.3386	-	-
4.2961	65000	0.3439	-	-
4.3622	66000	0.3378	-	-
4.4283	67000	0.2919	-	-
4.4944	68000	0.3099	-	-
4.5605	69000	0.2911	-	-
4.6266	70000	0.2644	-	-
4.6927	71000	0.3037	-	-
4.7588	72000	0.2862	-	-
4.8249	73000	0.2931	-	-
4.8909	74000	0.3613	-	-
4.9570	75000	0.3131	-	-

The bold row denotes the saved checkpoint.

Framework Versions

Python: 3.12.12
Sentence Transformers: 5.3.0
PyLate: 1.5.0
Transformers: 5.3.0
PyTorch: 2.9.0+cu130
Accelerate: 1.13.0
Datasets: 4.8.5
Tokenizers: 0.22.2

Citation

BibTeX

Sentence Transformers

bibtex

@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

bibtex

@inproceedings{DBLP:conf/cikm/ChaffinS25,
  author       = {Antoine Chaffin and
                  Rapha{"{e}}l Sourty},
  editor       = {Meeyoung Cha and
                  Chanyoung Park and
                  Noseong Park and
                  Carl Yang and
                  Senjuti Basu Roy and
                  Jessie Li and
                  Jaap Kamps and
                  Kijung Shin and
                  Bryan Hooi and
                  Lifang He},
  title        = {PyLate: Flexible Training and Retrieval for Late Interaction Models},
  booktitle    = {Proceedings of the 34th {ACM} International Conference on Information
                  and Knowledge Management, {CIKM} 2025, Seoul, Republic of Korea, November
                  10-14, 2025},
  pages        = {6334--6339},
  publisher    = {{ACM}},
  year         = {2025},
  url          = {https://github.com/lightonai/pylate},
  doi          = {10.1145/3746252.3761608},
}

	anchor	positive	negative_1
type	string	string	string
details	min: 9 tokens mean: 29.49 tokens max: 73 tokens	min: 8 tokens mean: 29.18 tokens max: 58 tokens	min: 4 tokens mean: 23.79 tokens max: 52 tokens