xtr-replicability/bge_small_xtr_contrastive_k512

PyLate model based on BAAI/bge-small-en-v1.5

This is a PyLate model finetuned from BAAI/bge-small-en-v1.5 on the msmarco-10m-triplets 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: BAAI/bge-small-en-v1.5
• Document Length: 300 tokens
• Query Length: 32 tokens
• Output Dimensionality: 128 tokens
• Similarity Function: MaxSim
• Training Dataset:
  • msmarco-10m-triplets

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': 300, 'do_lower_case': True, 'architecture': 'BertModel'})
  (1): Dense({'in_features': 384, '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="pylate_model_id",
)

# 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="pylate_model_id",
)

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

• Evaluated with pylate.evaluation.colbert_triplet.ColBERTTripletEvaluator

Metric	Value
accuracy	0.989

Training Details

Training Dataset

msmarco-10m-triplets

• Dataset: msmarco-10m-triplets at 8c5139a
• Size: 9,998,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: 32 tokens mean: 32.0 tokens max: 32 tokens	min: 32 tokens mean: 32.0 tokens max: 32 tokens	min: 32 tokens mean: 32.0 tokens max: 32 tokens

• Samples:

  query	positive	negative
  what is the intent of the encomienda system	Encomienda system established. The encomienda system was created by the Spanish to control and regulate American Indian labor and behavior during the colonization of the Americas. Under the encomienda system, conquistadors and other leaders (encomenderos) received grants of a number of Indians, from whom they could exact “tribute” in the form of gold or labor.	SearchNetworking. 1 Intent-based networking needed to run more complex networks Experts address why intent-based networking systems are needed to manage networks of the future that connect data center, public ...
  current weather in el cajon ca	El Cajon, CA Weather Forecast. TODAY - Patchy fog in the morning. Cloudy with a chance of rain. Highs around 63 in the western valleys to 56 to 61 near the foothills. Areas of winds southwest 15 mph. Chance of measurable precipitation 50 percent.	The two weather patterns alternate every five years and affect the weather of the entire continent. Summers in North America during the El Nino niño weather pattern are wetter than. Average winters In North america During El nino niño are warmer on average and there is less. snow falll Nino niño causes higher ocean water temperatures While La nina niña causes colder. Water temperatures the two weather patterns alternate every five years and affect the weather of the. Entire continent Summers In north america During The el nino niño weather pattern are. wetter than average
  what is meant by gathering natural gas	Our Gathering and Processing segment consists of gathering, compressing, dehydrating, treating, conditioning, processing, and marketing natural gas and gathering crude oil. The gathering of natural gas consists of aggregating natural gas produced from various wells through small diameter gathering lines to processing plants. Natural gas has a widely varying composition depending on the field, the formation and the reservoir from which it is produced.	Propane is a byproduct of the oil refining process and therefore fluctuates in price more so than natural gas. Since the USA is the Saudi Arabia of natural gas, it should be cheaper per BTU than propane.

• Loss: pylate.losses.xtr_primeqa.XTRPrimeQA

Evaluation Dataset

msmarco-10m-triplets

• Dataset: msmarco-10m-triplets at 8c5139a
• Size: 2,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: 32 tokens mean: 32.0 tokens max: 32 tokens	min: 32 tokens mean: 32.0 tokens max: 32 tokens	min: 32 tokens mean: 32.0 tokens max: 32 tokens

• Samples:

  query	positive	negative
  vegetables good for eye blood flow	A diet high in saturated fats and sugar lacks the antioxidants needed for good eye health. A diet high in saturated fats also creates substances that put the health of your eyes at risk, such as arterial plaque, which can cause restricted blood flow in the blood vessels in your eyes.ome Excellent Choices for Feeding Healthy Eyes. The following list of foods provides good to excellent antioxidant support for your eyes: 1 Green vegetables like broccoli, green beans, spinach, peas and Brussels sprouts. 2 Carrots, celery and corn. 3 Leafy lettuce. 4 Sweet potatoes and yams.	For this reason, a bag of frozen vegetables should be wrapped in a towel if used on extremities with very low blood flow, such as toes or hands. If the injury occurs in an area with little fat or muscle beneath the skin, such as fingers, take the compress off after 10 minutes maximum, wait 5 minutes and reapply.
  is battery fully charged	In both Living on 12 Volts with Ample Power, and Wiring 12 Volts for Ample Power the authors explain that a battery is fully charged when the voltage is about 14.4 Volts and current through the battery has declined to less than 2% of the capacity of the battery in Amp-hours ...2 Amps for a 100 Ah battery.	1 Do not fully charge or fully discharge your device’s battery — charge it to around 50%. 2 If you store a device when its battery is fully discharged, the battery could fall into a deep discharge state, which renders it incapable of holding a charge.
  what does the name adalina represent	Adalina /ada-li-na/ [4 sylls.] as a girls' name is of Old German derivation, and the name Adalina means noble. Adalina is a version of Adeline (Old German): French version of Adela. Andolina is a popular surname. Kreatif forms: Adalipa, Adalira, Adilina.	However, licensing as a broker or salesperson authorizes the licensee to represent parties on either side of a transaction. The choice of which side to represent is a business decision for the licensee. In the U.S., real estate brokers and salespersons are licensed by each state, not by the federal government.

• Loss: pylate.losses.xtr_primeqa.XTRPrimeQA

Training Hyperparameters

Non-Default Hyperparameters

• eval_strategy: steps
• per_device_train_batch_size: 196
• per_device_eval_batch_size: 196
• learning_rate: 3e-05
• max_grad_norm: 10.0
• num_train_epochs: 0
• max_steps: 50000
• warmup_ratio: 0.01
• bf16: True
• torch_compile: True
• torch_compile_backend: inductor
• eval_on_start: True

All Hyperparameters

Click to expand

• overwrite_output_dir: False
• do_predict: False
• eval_strategy: steps
• prediction_loss_only: True
• per_device_train_batch_size: 196
• per_device_eval_batch_size: 196
• 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: 3e-05
• weight_decay: 0.0
• adam_beta1: 0.9
• adam_beta2: 0.999
• adam_epsilon: 1e-08
• max_grad_norm: 10.0
• num_train_epochs: 0
• max_steps: 50000
• lr_scheduler_type: linear
• lr_scheduler_kwargs: {}
• warmup_ratio: 0.01
• 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: 0
• ddp_backend: None
• tpu_num_cores: None
• tpu_metrics_debug: False
• debug: []
• dataloader_drop_last: False
• dataloader_num_workers: 0
• 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: True
• 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: True
• torch_compile_backend: inductor
• 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: True
• use_liger_kernel: False
• liger_kernel_config: None
• eval_use_gather_object: False
• average_tokens_across_devices: False
• prompts: None
• batch_sampler: batch_sampler
• multi_dataset_batch_sampler: proportional
• router_mapping: {}
• learning_rate_mapping: {}