No official torrent yet — requesting it tells the flock what to host next. You cannot seed a catalog-only page.
Cataloged from lightonai/Agent-ModernColBERT
All models by lightonai →
A few weeks ago, we evaluated Reason-ModernColBERT, a 150M late-interaction model trained on ReasonIR data nearly solved BrowseComp-Plus, reaching 87.56% accuracy with GPT-5 (a +7.59 absolute jump over the previous SOTA) while topping recall and calibration error, while not being trained for agentic retrieval at all (and being one year old). We now present Agent-ModernColBERT, a model specifically fine-tuned for agentic retrieval using the AgentIR dataset released alongside AgentIR. You can find the training boilerplate here. This very lightweight fine-tuning adds increase the performance of Reason-ModernColBERT by another 10%, which allows, when exposing the get_document function and the GPT-OSS-120B model, to beat the original GPT-5 + Qwen3-8B runs, while using a retriever model 54× smaller and an open source LLM.
Before issuing a query, deep research agents generate explicit reasoning traces describing what they're looking for and why. Conventional retrievers throw all of that away. The AgentIR paper introduced Reasoning-Aware Retrieval: instead of discarding those reasoning traces, concatenate them to the query you send to the retriever. To enable this, the authors generated trajectories from a deep-research agent and released them (the Tevatron/AgentIR-data dataset), then trained AgentIR-4B, a 4B dense model that beats much larger baselines including ReasonIR-8B and rerank pipelines.
The get_document column indicates whether the LLM was exposed to the possibility of using the get_document function that allows, in addition to the usual search function that returns a snippet of the top-5 documents for a given query (the beginning of the document), to request the full content of a document using its id.
It allows the model to get more information when needed, but can also add noise to the process if the model reads full documents that are irrelevant, and thus require a precise retrieval signal.
A few things worth highlighting:
get_document function exposed, Agent-ModernColBERT + GPT-OSS-120B reaches 72.53, higher than the original GPT-5 + Qwen3-8B setup (which doesn't expose get_document).
Agent-ModernColBERT was trained on the AgentIR query format, which prepends an instruction and concatenates the agent's reasoning trace with the search query. Documents are encoded as plain text, no prefix.
The exact query layout is:
A few non-obvious details worth flagging:
\nQuery: (a backslash followed by n, 2 characters — not a newline). This is inherited byte-for-byte from the AgentIR training script and the model expects it.Query: appears twice in every query — once at the end of the prefix, once as a real label before the search query."Empty" (this matches the AgentIR agent loop).Helper:
Then encode normally with is_query=True for queries and is_query=False for documents:
If you rather want to use a simpler version without the prefix (just reasoning + query), you can use the no-prefix version, at the cost of slighlty lower performance:
get_document setupget_document is a net improvement in most cases, but the gains depend on both the retrieval model and the agent. The very strong numbers above are not a given for every combination. We tried adding more get_document data points, but were not able to reproduce the original AgentIR results exactly (small deltas remain after fixing environments). We therefore omit those points from the comparison: although they look favorable to Agent-ModernColBERT, we can't yet rule out GPT-OSS-side errors.
This is a multi-vector (ColBERT-style late interaction) embedding model finetuned from lightonai/GTE-ModernColBERT-v1 on the agent_ir-data dataset. It maps sentences & paragraphs to sequences of 128-dimensional dense vectors and can be used for semantic textual similarity using the MaxSim operator.
This model can be used with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via the MultiVectorEncoder:
First install the PyLate library:
Use this model with PyLate to index and retrieve documents. The index uses FastPLAID for efficient similarity search.
Load the ColBERT model and initialize the PLAID index, then encode and index your documents:
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:
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.
Note that queries must be wrapped with format_query() — see the Query format section above. Pass the agent's reasoning trace as the first argument, or None if you don't have one.
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 the rank function and pass the queries and documents to rerank.
As with retrieval, queries must be wrapped with format_query() (see the Query format section). Documents are encoded as plain text — no prefix.
query, positive, negative_0, negative_1, negative_2, negative_3, negative_4, negative_5, and negative_6pylate.losses.cached_contrastive.CachedContrastiveper_device_train_batch_size: 32per_device_eval_batch_size: 32learning_rate: 3e-06num_train_epochs: 2bf16: Truedataloader_num_workers: 8accelerator_config: {'split_batches': True, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}overwrite_output_dir: Falsedo_predict: Falseeval_strategy: noprediction_loss_only: Trueper_device_train_batch_size: 32per_device_eval_batch_size: 32per_gpu_train_batch_size: Noneper_gpu_eval_batch_size: Nonegradient_accumulation_steps: 1eval_accumulation_steps: Nonetorch_empty_cache_steps: Nonelearning_rate: 3e-06weight_decay: 0.0adam_beta1: 0.9adam_beta2: 0.999adam_epsilon: 1e-08max_grad_norm: 1.0num_train_epochs: 2max_steps: -1lr_scheduler_type: linearlr_scheduler_kwargs: {}warmup_ratio: 0.0warmup_steps: 0log_level: passivelog_level_replica: warninglog_on_each_node: Truelogging_nan_inf_filter: Truesave_safetensors: Truesave_on_each_node: Falsesave_only_model: Falserestore_callback_states_from_checkpoint: Falseno_cuda: Falseuse_cpu: Falseuse_mps_device: Falseseed: 42data_seed: Nonejit_mode_eval: Falsebf16: Truefp16: Falsefp16_opt_level: O1half_precision_backend: autobf16_full_eval: Falsefp16_full_eval: Falsetf32: Nonelocal_rank: 6ddp_backend: Nonetpu_num_cores: Nonetpu_metrics_debug: Falsedebug: []dataloader_drop_last: Truedataloader_num_workers: 8dataloader_prefetch_factor: Nonepast_index: -1disable_tqdm: Falseremove_unused_columns: Truelabel_names: Noneload_best_model_at_end: Falseignore_data_skip: Falsefsdp: []fsdp_min_num_params: 0fsdp_config: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}fsdp_transformer_layer_cls_to_wrap: Noneaccelerator_config: {'split_batches': True, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}parallelism_config: Nonedeepspeed: Nonelabel_smoothing_factor: 0.0optim: adamw_torch_fusedoptim_args: Noneadafactor: Falsegroup_by_length: Falselength_column_name: lengthproject: huggingfacetrackio_space_id: trackioddp_find_unused_parameters: Noneddp_bucket_cap_mb: Noneddp_broadcast_buffers: Falsedataloader_pin_memory: Truedataloader_persistent_workers: Falseskip_memory_metrics: Trueuse_legacy_prediction_loop: Falsepush_to_hub: Falseresume_from_checkpoint: Nonehub_model_id: Nonehub_strategy: every_savehub_private_repo: Nonehub_always_push: Falsehub_revision: Nonegradient_checkpointing: Falsegradient_checkpointing_kwargs: Noneinclude_inputs_for_metrics: Falseinclude_for_metrics: []eval_do_concat_batches: Truefp16_backend: autopush_to_hub_model_id: Nonepush_to_hub_organization: Nonemp_parameters:auto_find_batch_size: Falsefull_determinism: Falsetorchdynamo: Noneray_scope: lastddp_timeout: 1800torch_compile: Falsetorch_compile_backend: Nonetorch_compile_mode: Noneinclude_tokens_per_second: Falseinclude_num_input_tokens_seen: noneftune_noise_alpha: Noneoptim_target_modules: Nonebatch_eval_metrics: Falseeval_on_start: Falseuse_liger_kernel: Falseliger_kernel_config: Noneeval_use_gather_object: Falseaverage_tokens_across_devices: Trueprompts: Nonebatch_sampler: batch_samplermulti_dataset_batch_sampler: proportionalrouter_mapping: {}learning_rate_mapping: {}<INSTRUCT_PREFIX>Reasoning: <reasoning>\n\nQuery: <search_query>INSTRUCT_PREFIX = ( "Instruct: Given a user's reasoning followed by a web search query, " "retrieve relevant passages that answer the query while incorporating the user's reasoning" "\\nQuery:" # literal backslash-n, NOT a newline — keep as-is)
def format_query(reasoning: str | None, search_query: str) -> str: r = reasoning if reasoning else "Empty" return f"{INSTRUCT_PREFIX}Reasoning: {r}\n\nQuery: {search_query}"from pylate import models
model = models.ColBERT(model_name_or_path="lightonai/Agent-ModernColBERT")
queries = [ format_query( reasoning="The user wants a code-level explanation of PLAID, so passages " "with implementation details should be prioritized.", search_query="What does the PLAID centroid scoring step do?", ),]documents = [ "PLAID first scores documents against quantized centroids of the ColBERT " "token embeddings, then re-ranks a short candidate list with the full " "late-interaction score.",]
query_embs = model.encode(queries, is_query=True)doc_embs = model.encode(documents, is_query=False)ColBERT( (0): Transformer({'max_seq_length': 8191, '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}))pip install "sentence-transformers>=6.0.0"from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder("lightonai/Agent-ModernColBERT")
query = "Which planet is known as the Red Planet?"documents = [ "Venus is often called Earth's twin because of its similar size and proximity.", "Mars, known for its reddish appearance, is often referred to as the Red Planet.", "Jupiter, the largest planet in our solar system, has a prominent red spot.", "Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",]
query_embeddings = model.encode_query(query)document_embeddings = model.encode_document(documents)print(query_embeddings.shape, document_embeddings[0].shape)# (12, 128) (18, 128)
# MaxSim late-interaction scoring (higher is more relevant)scores = model.similarity(query_embeddings, document_embeddings)print(scores)# tensor([[10.4239, 11.1039, 10.6760, 10.9058]])pip install -U pylatefrom pylate import indexes, models, retrieve
# Step 1: Load the ColBERT modelmodel = models.ColBERT( model_name_or_path="pylate_model_id",)
# Step 2: Initialize the PLAID indexindex = indexes.PLAID( index_folder="pylate-index", index_name="index", override=True, # This overwrites the existing index if any)
# Step 3: Encode the documentsdocuments_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 idsindex.add_documents( documents_ids=documents_ids, documents_embeddings=documents_embeddings,)# To load an index, simply instantiate it with the correct folder/name and without overriding itindex = indexes.PLAID( index_folder="pylate-index", index_name="index",)INSTRUCT_PREFIX = ( "Instruct: Given a user's reasoning followed by a web search query, " "retrieve relevant passages that answer the query while incorporating the user's reasoning" "\\nQuery:" # literal backslash-n, NOT a newline — keep as-is)
def format_query(reasoning: str | None, search_query: str) -> str: r = reasoning if reasoning else "Empty" return f"{INSTRUCT_PREFIX}Reasoning: {r}\n\nQuery: {search_query}"# Step 1: Initialize the ColBERT retrieverretriever = retrieve.ColBERT(index=index)
# Step 2: Build queries in the AgentIR format, then encode themraw_queries = [ ( "The user wants a code-level explanation of PLAID, so passages with " "implementation details or pseudo-code should be prioritized.", "What does the PLAID centroid scoring step do?", ), ( "The user is comparing late-interaction models to dense bi-encoders, " "likely with a focus on latency trade-offs.", "How does ColBERT compare to dense retrievers in terms of latency?", ),]queries = [format_query(reasoning, q) for reasoning, q in raw_queries]
queries_embeddings = model.encode( queries, batch_size=32, is_query=True, # adds the [Q] marker + uses query_length show_progress_bar=True,)
# Step 3: Retrieve top-k documentsscores = retriever.retrieve( queries_embeddings=queries_embeddings, k=10, # Retrieve the top 10 matches for each query)from pylate import rank, modelsINSTRUCT_PREFIX = ( "Instruct: Given a user's reasoning followed by a web search query, " "retrieve relevant passages that answer the query while incorporating the user's reasoning" "\\nQuery:" # literal backslash-n, NOT a newline — keep as-is)
def format_query(reasoning: str | None, search_query: str) -> str: r = reasoning if reasoning else "Empty" return f"{INSTRUCT_PREFIX}Reasoning: {r}\n\nQuery: {search_query}"
raw_queries = [ ( "The user wants to understand how MaxSim scoring works at a high level.", "What is late interaction in ColBERT?", ), ( "The user is benchmarking retrievers and cares about index size.", "How large is a PLAID index compared to a dense FAISS index?", ),]queries = [format_query(reasoning, q) for reasoning, q in raw_queries]
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="pylate_model_id",)
queries_embeddings = model.encode( queries, is_query=True, # adds the [Q] marker + uses query_length)
documents_embeddings = model.encode( documents, is_query=False, # adds the [D] marker + uses document_length)
reranked_documents = rank.rerank( documents_ids=documents_ids, queries_embeddings=queries_embeddings, documents_embeddings=documents_embeddings,)@misc{Agent-ModernColBERT,title = {Agent-ModernColBERT},author = {Chaffin, Antoine},url = {https://huggingface.co/lightonai/Agent-ModernColBERT},year = {2026}}@article{DBLP:journals/corr/abs-2603-04384, author = {Zijian Chen and Xueguang Ma and Shengyao Zhuang and Jimmy Lin and Akari Asai and Victor Zhong}, title = {AgentIR: Reasoning-Aware Retrieval for Deep Research Agents}, journal = {CoRR}, volume = {abs/2603.04384}, year = {2026}, url = {https://doi.org/10.48550/arXiv.2603.04384}, doi = {10.48550/ARXIV.2603.04384}, eprinttype = {arXiv}, eprint = {2603.04384}, timestamp = {Wed, 08 Apr 2026 10:55:23 +0200}, biburl = {https://dblp.org/rec/journals/corr/abs-2603-04384.bib}, bibsource = {dblp computer science bibliography, https://dblp.org}}@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"}@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},}@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}}