Add nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise

Upload fine-tuned rerankers for BioASQ 14B

Co-authored-by: André Ribeiro <andrepedro2004@hotmail.com>
Co-authored-by: Rúben Garrido <rubengarrido@ua.pt>

.gitattributes

@@ -36,3 +36,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 BAAI-bge-reranker-base-E2-S1-Mpairwise-FullDataTrue/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 BAAI-bge-reranker-v2-m3-E2-S1-Mpairwise-FullDataTrue/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 nvidia-llama-nemotron-rerank-1b-v2-E2-S4-Mmulti_neg_pairwise-Linfonce-FullData/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/tokenizer.json filter=lfs diff=lfs merge=lfs -text

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/config.json

@@ -0,0 +1,48 @@
+{
+  "architectures": [
+    "LlamaBidirectionalForSequenceClassification"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "llama_bidirectional_model.LlamaBidirectionalConfig",
+    "AutoModelForSequenceClassification": "llama_bidirectional_model.LlamaBidirectionalForSequenceClassification"
+  },
+  "bos_token_id": 128000,
+  "dtype": "bfloat16",
+  "eos_token_id": 128001,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "id2label": {
+    "0": "LABEL_0"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "label2id": {
+    "LABEL_0": 0
+  },
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama_bidirec",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "pad_token_id": 128001,
+  "pooling": "avg",
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_theta": 500000.0,
+    "rope_type": "llama3"
+  },
+  "temperature": 1.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.2.0",
+  "use_cache": false,
+  "vocab_size": 128256
+}

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/llama_bidirectional_model.py

@@ -0,0 +1,362 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0.
+"""
+Bidirectional Llama model for cross-encoder reranking.
+
+Modifies LlamaModel to use bidirectional (non-causal) attention so each token
+attends to all others — required for cross-encoder scoring of query-document pairs.
+
+Provides three classes:
+    - LlamaBidirectionalConfig: Adds pooling and temperature to LlamaConfig.
+    - LlamaBidirectionalModel: LlamaModel with causal masking replaced by
+      bidirectional masking. Overrides forward() to support transformers >=4.44.
+    - LlamaBidirectionalForSequenceClassification: Pools hidden states and
+      projects to a relevance score via a linear head.
+
+Transformers version compatibility (>=4.44 including 5.0+):
+    The forward() implementation handles these API changes at import time via
+    inspect.signature() on LlamaDecoderLayer and DynamicCache:
+
+    < 4.53:  _update_causal_mask exists on LlamaModel (not used here).
+    4.53+:   Masking moved to masking_utils; requires full forward() override.
+    < 4.54:  Decoder layer returns a tuple.
+    4.54+:   Decoder layer returns a tensor.
+    < 4.56:  Cache kwarg is ``past_key_value`` (singular).
+    4.56+:   Cache kwarg is ``past_key_values`` (plural); DynamicCache accepts config.
+    5.0+:    Native ``create_bidirectional_mask`` in masking_utils.
+"""
+
+import inspect
+from typing import Optional, Union, Tuple, List
+
+import torch
+import torch.nn as nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.modeling_outputs import SequenceClassifierOutputWithPast
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import (
+    LlamaDecoderLayer,
+    LlamaModel,
+    LlamaPreTrainedModel,
+)
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+# Check if native create_bidirectional_mask exists (transformers >= 5.0)
+try:
+    from transformers.masking_utils import create_bidirectional_mask
+
+    _HAS_NATIVE_BIDIRECTIONAL_MASK = True
+except ImportError:
+    from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+
+    _HAS_NATIVE_BIDIRECTIONAL_MASK = False
+
+# Detect API differences via introspection
+_decoder_forward_params = inspect.signature(LlamaDecoderLayer.forward).parameters
+_dynamic_cache_init_params = inspect.signature(DynamicCache.__init__).parameters
+
+# past_key_value (singular) in < 4.56, past_key_values (plural) in >= 4.56
+_USE_PLURAL_CACHE_PARAM = "past_key_values" in _decoder_forward_params
+# DynamicCache accepts config parameter in >= 4.56
+_DYNAMIC_CACHE_ACCEPTS_CONFIG = "config" in _dynamic_cache_init_params
+
+
+class LlamaBidirectionalConfig(LlamaConfig):
+    """Configuration for LlamaBidirectionalModel with pooling and temperature settings."""
+
+    model_type = "llama_bidirec"
+
+    def __init__(
+        self, pooling: str = "avg", temperature: float = 1.0, **kwargs
+    ) -> None:
+        """
+        Initialize bidirectional Llama configuration.
+
+        Args:
+            pooling: Pooling strategy for embeddings ("avg", "cls", "last", etc.)
+            temperature: Temperature scaling for embeddings
+            **kwargs: Additional arguments passed to LlamaConfig
+        """
+        self.pooling = pooling
+        self.temperature = temperature
+        super().__init__(**kwargs)
+
+
+class LlamaBidirectionalModel(LlamaModel):
+    """
+    LlamaModel modified to use bidirectional (non-causal) attention.
+
+    In standard Llama, each token can only attend to previous tokens (causal attention).
+    This model removes that restriction, allowing each token to attend to all tokens
+    in the sequence, which is useful for embedding tasks.
+
+    The key modifications are:
+        1. Setting is_causal=False on all attention layers
+        2. Using a bidirectional attention mask instead of causal mask
+    """
+
+    config_class = LlamaBidirectionalConfig
+
+    def __init__(self, config: LlamaConfig) -> None:
+        super().__init__(config)
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+
+    def _create_bidirectional_mask(
+        self,
+        input_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor | None:
+        """
+        Create bidirectional attention mask.
+
+        Args:
+            input_embeds: Input embeddings tensor of shape (batch_size, seq_len, hidden_size)
+            attention_mask: Optional 2D attention mask of shape (batch_size, seq_len)
+                where 1 indicates tokens to attend to and 0 indicates masked tokens
+
+        Returns:
+            4D attention mask suitable for the attention implementation, or None
+            if no masking is needed
+        """
+        if attention_mask is None:
+            return None
+
+        if _HAS_NATIVE_BIDIRECTIONAL_MASK:
+            return create_bidirectional_mask(
+                config=self.config,
+                input_embeds=input_embeds,
+                attention_mask=attention_mask,
+            )
+
+        # Fallback for transformers < 5.0 without create_bidirectional_mask
+
+        # Flash attention handles 2D masks internally; only pass mask if there
+        # are actually masked tokens (zeros), otherwise return None for efficiency
+        if getattr(self.config, "_attn_implementation", None) == "flash_attention_2":
+            has_masked_tokens = (attention_mask == 0).any()
+            return attention_mask if has_masked_tokens else None
+
+        return _prepare_4d_attention_mask(attention_mask, input_embeds.dtype)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        """
+        Forward pass with bidirectional attention.
+
+        Args:
+            input_ids: Input token IDs of shape (batch_size, seq_len)
+            attention_mask: Attention mask of shape (batch_size, seq_len)
+            position_ids: Position IDs for rotary embeddings
+            past_key_values: Cached key/value states for incremental decoding
+            inputs_embeds: Pre-computed input embeddings (alternative to input_ids)
+            cache_position: Position indices for cache updates
+            use_cache: Whether to return cached key/value states
+            **kwargs: Additional arguments passed to decoder layers
+
+        Returns:
+            BaseModelOutputWithPast containing last_hidden_state and past_key_values
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # Initialize cache if needed
+        if use_cache and past_key_values is None:
+            if _DYNAMIC_CACHE_ACCEPTS_CONFIG:
+                past_key_values = DynamicCache(config=self.config)
+            else:
+                past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        bidirectional_mask = self._create_bidirectional_mask(
+            inputs_embeds, attention_mask
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # Build decoder layer kwargs with correct cache parameter name
+        # (past_key_value in < 4.56, past_key_values in >= 4.56)
+        layer_kwargs = {
+            "attention_mask": bidirectional_mask,
+            "position_ids": position_ids,
+            "use_cache": use_cache,
+            "cache_position": cache_position,
+            "position_embeddings": position_embeddings,
+        }
+        if _USE_PLURAL_CACHE_PARAM:
+            layer_kwargs["past_key_values"] = past_key_values
+        else:
+            layer_kwargs["past_key_value"] = past_key_values
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            layer_outputs = decoder_layer(hidden_states, **layer_kwargs)
+
+            # Decoder returns tuple in < 4.54, tensor in >= 4.54
+            if isinstance(layer_outputs, tuple):
+                hidden_states = layer_outputs[0]
+            else:
+                hidden_states = layer_outputs
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+def pool(
+    last_hidden_states: torch.Tensor, attention_mask: torch.Tensor, pool_type: str
+) -> torch.Tensor:
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+
+    if pool_type == "avg":
+        emb = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+    elif pool_type == "weighted_avg":
+        emb = last_hidden.sum(dim=1)
+    elif pool_type == "cls":
+        emb = last_hidden[:, 0]
+    elif pool_type == "last":
+        left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
+        if left_padding:
+            emb = last_hidden[:, -1]
+        else:
+            sequence_lengths = attention_mask.sum(dim=1) - 1
+            batch_size = last_hidden.shape[0]
+            emb = last_hidden[
+                torch.arange(batch_size, device=last_hidden.device), sequence_lengths
+            ]
+    else:
+        raise ValueError(f"pool_type {pool_type} not supported")
+
+    return emb
+
+
+class LlamaBidirectionalForSequenceClassification(LlamaPreTrainedModel):
+    config_class = LlamaBidirectionalConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+        self.model = LlamaBidirectionalModel(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        transformer_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs,
+        )
+        hidden_states = transformer_outputs[0]
+
+        pooled_hidden_states = pool(
+            last_hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            pool_type=self.config.pooling,
+        )
+
+        pooled_logits = self.score(pooled_hidden_states)
+        pooled_logits = pooled_logits / self.config.temperature
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(pooled_logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(
+                    pooled_logits.view(-1, self.num_labels), labels.view(-1)
+                )
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d7ff6ae8ef3212cebe4264b9de6cb71603f1f4cb1bd7e860631c94e38a44df8
+size 2471649792

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/ranx_results.json

@@ -0,0 +1,31 @@
+{
+    "model": "nvidia-llama-nemotron-rerank-1b-v2",
+    "val_files": [
+        "../../data/val_data/13B1_golden.json",
+        "../../data/val_data/13B2_golden.json"
+    ],
+    "total": {
+        "ndcg@5": 0.9978289985504203,
+        "mrr": 0.9970588235294118,
+        "recall@10": 0.6594908184294858,
+        "recall@100": 1.0,
+        "map@10": 0.6564445999420908,
+        "map-bioasq@10": 0.996953781512605
+    },
+    "13B1_golden.json": {
+        "ndcg@5": 0.9956579971008406,
+        "mrr": 0.9941176470588236,
+        "recall@10": 0.7100945834410882,
+        "recall@100": 1.0,
+        "map@10": 0.7040021464662983,
+        "map-bioasq@10": 0.9939075630252101
+    },
+    "13B2_golden.json": {
+        "ndcg@5": 1.0,
+        "mrr": 1.0,
+        "recall@10": 0.6088870534178833,
+        "recall@100": 1.0,
+        "map@10": 0.6088870534178833,
+        "map-bioasq@10": 1.0
+    }
+}

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c85066e7642934ed09b44155e6566b0b5dab2637fb9433439ba5c9c7f8b50d3
+size 17210018

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|end_of_text|>",
+  "is_local": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 4096,
+  "pad_token": "<|end_of_text|>",
+  "tokenizer_class": "TokenizersBackend"
+}

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/trainer_state.json

@@ -0,0 +1,757 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 2.0,
+  "eval_steps": 500,
+  "global_step": 5098,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.019615535504119264,
+      "grad_norm": 0.0,
+      "learning_rate": 9.903883876029816e-05,
+      "loss": 0.20531219482421875,
+      "step": 50
+    },
+    {
+      "epoch": 0.03923107100823853,
+      "grad_norm": 0.0,
+      "learning_rate": 9.805806198509219e-05,
+      "loss": 0.008828125,
+      "step": 100
+    },
+    {
+      "epoch": 0.05884660651235779,
+      "grad_norm": 6.46875,
+      "learning_rate": 9.707728520988623e-05,
+      "loss": 0.3015625,
+      "step": 150
+    },
+    {
+      "epoch": 0.07846214201647705,
+      "grad_norm": 13.875,
+      "learning_rate": 9.609650843468027e-05,
+      "loss": 0.330703125,
+      "step": 200
+    },
+    {
+      "epoch": 0.09807767752059632,
+      "grad_norm": 0.0,
+      "learning_rate": 9.51157316594743e-05,
+      "loss": 0.0034375,
+      "step": 250
+    },
+    {
+      "epoch": 0.11769321302471558,
+      "grad_norm": 0.0,
+      "learning_rate": 9.413495488426834e-05,
+      "loss": 0.0003125,
+      "step": 300
+    },
+    {
+      "epoch": 0.13730874852883485,
+      "grad_norm": 0.0,
+      "learning_rate": 9.315417810906238e-05,
+      "loss": 0.0015625,
+      "step": 350
+    },
+    {
+      "epoch": 0.1569242840329541,
+      "grad_norm": 0.0,
+      "learning_rate": 9.217340133385642e-05,
+      "loss": 0.00125,
+      "step": 400
+    },
+    {
+      "epoch": 0.17653981953707337,
+      "grad_norm": 0.0,
+      "learning_rate": 9.119262455865045e-05,
+      "loss": 0.000625,
+      "step": 450
+    },
+    {
+      "epoch": 0.19615535504119264,
+      "grad_norm": 0.0,
+      "learning_rate": 9.021184778344449e-05,
+      "loss": 0.0103125,
+      "step": 500
+    },
+    {
+      "epoch": 0.2157708905453119,
+      "grad_norm": 0.0,
+      "learning_rate": 8.923107100823853e-05,
+      "loss": 0.0,
+      "step": 550
+    },
+    {
+      "epoch": 0.23538642604943116,
+      "grad_norm": 0.0,
+      "learning_rate": 8.825029423303257e-05,
+      "loss": 0.0,
+      "step": 600
+    },
+    {
+      "epoch": 0.2550019615535504,
+      "grad_norm": 0.0,
+      "learning_rate": 8.72695174578266e-05,
+      "loss": 0.0,
+      "step": 650
+    },
+    {
+      "epoch": 0.2746174970576697,
+      "grad_norm": 0.0,
+      "learning_rate": 8.628874068262064e-05,
+      "loss": 0.0,
+      "step": 700
+    },
+    {
+      "epoch": 0.29423303256178895,
+      "grad_norm": 0.0,
+      "learning_rate": 8.530796390741468e-05,
+      "loss": 0.0,
+      "step": 750
+    },
+    {
+      "epoch": 0.3138485680659082,
+      "grad_norm": 0.0,
+      "learning_rate": 8.432718713220872e-05,
+      "loss": 0.0,
+      "step": 800
+    },
+    {
+      "epoch": 0.3334641035700275,
+      "grad_norm": 0.0,
+      "learning_rate": 8.334641035700274e-05,
+      "loss": 0.0,
+      "step": 850
+    },
+    {
+      "epoch": 0.35307963907414674,
+      "grad_norm": 0.0,
+      "learning_rate": 8.236563358179678e-05,
+      "loss": 0.0034375,
+      "step": 900
+    },
+    {
+      "epoch": 0.372695174578266,
+      "grad_norm": 0.0,
+      "learning_rate": 8.138485680659082e-05,
+      "loss": 0.0,
+      "step": 950
+    },
+    {
+      "epoch": 0.39231071008238527,
+      "grad_norm": 0.0,
+      "learning_rate": 8.040408003138486e-05,
+      "loss": 0.0003125,
+      "step": 1000
+    },
+    {
+      "epoch": 0.41192624558650454,
+      "grad_norm": 0.0,
+      "learning_rate": 7.942330325617889e-05,
+      "loss": 0.0,
+      "step": 1050
+    },
+    {
+      "epoch": 0.4315417810906238,
+      "grad_norm": 0.0,
+      "learning_rate": 7.844252648097293e-05,
+      "loss": 0.0,
+      "step": 1100
+    },
+    {
+      "epoch": 0.45115731659474306,
+      "grad_norm": 0.0,
+      "learning_rate": 7.746174970576697e-05,
+      "loss": 0.0,
+      "step": 1150
+    },
+    {
+      "epoch": 0.4707728520988623,
+      "grad_norm": 0.0,
+      "learning_rate": 7.648097293056101e-05,
+      "loss": 0.0028125,
+      "step": 1200
+    },
+    {
+      "epoch": 0.49038838760298153,
+      "grad_norm": 0.0,
+      "learning_rate": 7.550019615535504e-05,
+      "loss": 0.0,
+      "step": 1250
+    },
+    {
+      "epoch": 0.5100039231071009,
+      "grad_norm": 0.0,
+      "learning_rate": 7.451941938014908e-05,
+      "loss": 0.0,
+      "step": 1300
+    },
+    {
+      "epoch": 0.5296194586112201,
+      "grad_norm": 0.0,
+      "learning_rate": 7.353864260494312e-05,
+      "loss": 0.0,
+      "step": 1350
+    },
+    {
+      "epoch": 0.5492349941153394,
+      "grad_norm": 0.0,
+      "learning_rate": 7.255786582973716e-05,
+      "loss": 0.0,
+      "step": 1400
+    },
+    {
+      "epoch": 0.5688505296194586,
+      "grad_norm": 0.0,
+      "learning_rate": 7.157708905453119e-05,
+      "loss": 0.0,
+      "step": 1450
+    },
+    {
+      "epoch": 0.5884660651235779,
+      "grad_norm": 0.0,
+      "learning_rate": 7.059631227932523e-05,
+      "loss": 0.0,
+      "step": 1500
+    },
+    {
+      "epoch": 0.6080816006276971,
+      "grad_norm": 0.0,
+      "learning_rate": 6.961553550411927e-05,
+      "loss": 0.0,
+      "step": 1550
+    },
+    {
+      "epoch": 0.6276971361318164,
+      "grad_norm": 0.0,
+      "learning_rate": 6.863475872891331e-05,
+      "loss": 0.0,
+      "step": 1600
+    },
+    {
+      "epoch": 0.6473126716359356,
+      "grad_norm": 0.0,
+      "learning_rate": 6.765398195370734e-05,
+      "loss": 0.0,
+      "step": 1650
+    },
+    {
+      "epoch": 0.666928207140055,
+      "grad_norm": 0.0,
+      "learning_rate": 6.667320517850138e-05,
+      "loss": 0.0,
+      "step": 1700
+    },
+    {
+      "epoch": 0.6865437426441742,
+      "grad_norm": 0.0,
+      "learning_rate": 6.569242840329542e-05,
+      "loss": 0.0,
+      "step": 1750
+    },
+    {
+      "epoch": 0.7061592781482935,
+      "grad_norm": 0.0,
+      "learning_rate": 6.471165162808946e-05,
+      "loss": 0.0,
+      "step": 1800
+    },
+    {
+      "epoch": 0.7257748136524127,
+      "grad_norm": 0.0,
+      "learning_rate": 6.373087485288348e-05,
+      "loss": 0.0,
+      "step": 1850
+    },
+    {
+      "epoch": 0.745390349156532,
+      "grad_norm": 0.0,
+      "learning_rate": 6.275009807767752e-05,
+      "loss": 0.004375,
+      "step": 1900
+    },
+    {
+      "epoch": 0.7650058846606512,
+      "grad_norm": 0.0,
+      "learning_rate": 6.176932130247156e-05,
+      "loss": 0.0009375,
+      "step": 1950
+    },
+    {
+      "epoch": 0.7846214201647705,
+      "grad_norm": 0.0,
+      "learning_rate": 6.07885445272656e-05,
+      "loss": 0.0,
+      "step": 2000
+    },
+    {
+      "epoch": 0.8042369556688898,
+      "grad_norm": 0.0,
+      "learning_rate": 5.980776775205963e-05,
+      "loss": 0.0,
+      "step": 2050
+    },
+    {
+      "epoch": 0.8238524911730091,
+      "grad_norm": 0.0,
+      "learning_rate": 5.8826990976853666e-05,
+      "loss": 0.00078125,
+      "step": 2100
+    },
+    {
+      "epoch": 0.8434680266771283,
+      "grad_norm": 0.0,
+      "learning_rate": 5.7846214201647706e-05,
+      "loss": 0.0,
+      "step": 2150
+    },
+    {
+      "epoch": 0.8630835621812476,
+      "grad_norm": 0.0,
+      "learning_rate": 5.6865437426441747e-05,
+      "loss": 0.0,
+      "step": 2200
+    },
+    {
+      "epoch": 0.8826990976853668,
+      "grad_norm": 0.0,
+      "learning_rate": 5.5884660651235773e-05,
+      "loss": 0.0,
+      "step": 2250
+    },
+    {
+      "epoch": 0.9023146331894861,
+      "grad_norm": 0.0,
+      "learning_rate": 5.4903883876029814e-05,
+      "loss": 0.0,
+      "step": 2300
+    },
+    {
+      "epoch": 0.9219301686936053,
+      "grad_norm": 0.0,
+      "learning_rate": 5.3923107100823854e-05,
+      "loss": 0.0,
+      "step": 2350
+    },
+    {
+      "epoch": 0.9415457041977247,
+      "grad_norm": 0.0,
+      "learning_rate": 5.2942330325617895e-05,
+      "loss": 0.00125,
+      "step": 2400
+    },
+    {
+      "epoch": 0.9611612397018439,
+      "grad_norm": 0.0,
+      "learning_rate": 5.196155355041192e-05,
+      "loss": 0.0,
+      "step": 2450
+    },
+    {
+      "epoch": 0.9807767752059631,
+      "grad_norm": 0.0,
+      "learning_rate": 5.098077677520596e-05,
+      "loss": 0.0,
+      "step": 2500
+    },
+    {
+      "epoch": 1.0,
+      "eval_loss": 0.0,
+      "eval_runtime": 4.9531,
+      "eval_samples_per_second": 33.918,
+      "eval_steps_per_second": 1.211,
+      "step": 2549
+    },
+    {
+      "epoch": 1.0003923107100823,
+      "grad_norm": 0.0,
+      "learning_rate": 5e-05,
+      "loss": 0.0,
+      "step": 2550
+    },
+    {
+      "epoch": 1.0200078462142017,
+      "grad_norm": 0.0,
+      "learning_rate": 4.901922322479404e-05,
+      "loss": 0.0,
+      "step": 2600
+    },
+    {
+      "epoch": 1.039623381718321,
+      "grad_norm": 0.0,
+      "learning_rate": 4.8038446449588076e-05,
+      "loss": 0.0,
+      "step": 2650
+    },
+    {
+      "epoch": 1.0592389172224401,
+      "grad_norm": 0.0,
+      "learning_rate": 4.705766967438212e-05,
+      "loss": 0.0,
+      "step": 2700
+    },
+    {
+      "epoch": 1.0788544527265593,
+      "grad_norm": 0.0,
+      "learning_rate": 4.607689289917615e-05,
+      "loss": 0.0,
+      "step": 2750
+    },
+    {
+      "epoch": 1.0984699882306788,
+      "grad_norm": 0.0,
+      "learning_rate": 4.509611612397019e-05,
+      "loss": 0.0,
+      "step": 2800
+    },
+    {
+      "epoch": 1.118085523734798,
+      "grad_norm": 0.0,
+      "learning_rate": 4.4115339348764225e-05,
+      "loss": 0.0,
+      "step": 2850
+    },
+    {
+      "epoch": 1.1377010592389172,
+      "grad_norm": 0.0,
+      "learning_rate": 4.3134562573558265e-05,
+      "loss": 0.0,
+      "step": 2900
+    },
+    {
+      "epoch": 1.1573165947430364,
+      "grad_norm": 0.0,
+      "learning_rate": 4.21537857983523e-05,
+      "loss": 0.0,
+      "step": 2950
+    },
+    {
+      "epoch": 1.1769321302471558,
+      "grad_norm": 0.0,
+      "learning_rate": 4.117300902314634e-05,
+      "loss": 0.0,
+      "step": 3000
+    },
+    {
+      "epoch": 1.196547665751275,
+      "grad_norm": 0.0,
+      "learning_rate": 4.019223224794037e-05,
+      "loss": 0.0,
+      "step": 3050
+    },
+    {
+      "epoch": 1.2161632012553942,
+      "grad_norm": 0.0,
+      "learning_rate": 3.9211455472734406e-05,
+      "loss": 0.0,
+      "step": 3100
+    },
+    {
+      "epoch": 1.2357787367595134,
+      "grad_norm": 0.0,
+      "learning_rate": 3.823067869752845e-05,
+      "loss": 0.0,
+      "step": 3150
+    },
+    {
+      "epoch": 1.2553942722636329,
+      "grad_norm": 0.0,
+      "learning_rate": 3.724990192232248e-05,
+      "loss": 0.0,
+      "step": 3200
+    },
+    {
+      "epoch": 1.275009807767752,
+      "grad_norm": 0.0,
+      "learning_rate": 3.6269125147116514e-05,
+      "loss": 0.00046875,
+      "step": 3250
+    },
+    {
+      "epoch": 1.2946253432718713,
+      "grad_norm": 0.0,
+      "learning_rate": 3.5288348371910554e-05,
+      "loss": 0.0,
+      "step": 3300
+    },
+    {
+      "epoch": 1.3142408787759905,
+      "grad_norm": 0.0,
+      "learning_rate": 3.430757159670459e-05,
+      "loss": 0.0,
+      "step": 3350
+    },
+    {
+      "epoch": 1.33385641428011,
+      "grad_norm": 0.0,
+      "learning_rate": 3.332679482149863e-05,
+      "loss": 0.0,
+      "step": 3400
+    },
+    {
+      "epoch": 1.3534719497842291,
+      "grad_norm": 0.0,
+      "learning_rate": 3.234601804629266e-05,
+      "loss": 0.0,
+      "step": 3450
+    },
+    {
+      "epoch": 1.3730874852883483,
+      "grad_norm": 0.0,
+      "learning_rate": 3.13652412710867e-05,
+      "loss": 0.0,
+      "step": 3500
+    },
+    {
+      "epoch": 1.3927030207924678,
+      "grad_norm": 0.0,
+      "learning_rate": 3.0384464495880736e-05,
+      "loss": 0.0,
+      "step": 3550
+    },
+    {
+      "epoch": 1.412318556296587,
+      "grad_norm": 0.0,
+      "learning_rate": 2.9403687720674777e-05,
+      "loss": 0.0,
+      "step": 3600
+    },
+    {
+      "epoch": 1.4319340918007062,
+      "grad_norm": 0.0,
+      "learning_rate": 2.842291094546881e-05,
+      "loss": 0.0,
+      "step": 3650
+    },
+    {
+      "epoch": 1.4515496273048254,
+      "grad_norm": 0.0,
+      "learning_rate": 2.744213417026285e-05,
+      "loss": 0.0,
+      "step": 3700
+    },
+    {
+      "epoch": 1.4711651628089446,
+      "grad_norm": 0.0,
+      "learning_rate": 2.6461357395056884e-05,
+      "loss": 0.0,
+      "step": 3750
+    },
+    {
+      "epoch": 1.4907806983130638,
+      "grad_norm": 0.0,
+      "learning_rate": 2.5480580619850925e-05,
+      "loss": 0.0,
+      "step": 3800
+    },
+    {
+      "epoch": 1.5103962338171832,
+      "grad_norm": 0.0,
+      "learning_rate": 2.4499803844644962e-05,
+      "loss": 0.0,
+      "step": 3850
+    },
+    {
+      "epoch": 1.5300117693213025,
+      "grad_norm": 0.0,
+      "learning_rate": 2.3519027069439e-05,
+      "loss": 0.0,
+      "step": 3900
+    },
+    {
+      "epoch": 1.5496273048254219,
+      "grad_norm": 0.0,
+      "learning_rate": 2.2538250294233036e-05,
+      "loss": 0.0,
+      "step": 3950
+    },
+    {
+      "epoch": 1.569242840329541,
+      "grad_norm": 0.0,
+      "learning_rate": 2.1557473519027073e-05,
+      "loss": 0.0,
+      "step": 4000
+    },
+    {
+      "epoch": 1.5888583758336603,
+      "grad_norm": 0.0,
+      "learning_rate": 2.057669674382111e-05,
+      "loss": 0.0,
+      "step": 4050
+    },
+    {
+      "epoch": 1.6084739113377795,
+      "grad_norm": 0.0,
+      "learning_rate": 1.9595919968615144e-05,
+      "loss": 0.0,
+      "step": 4100
+    },
+    {
+      "epoch": 1.6280894468418987,
+      "grad_norm": 0.0,
+      "learning_rate": 1.861514319340918e-05,
+      "loss": 0.0,
+      "step": 4150
+    },
+    {
+      "epoch": 1.647704982346018,
+      "grad_norm": 0.0,
+      "learning_rate": 1.7634366418203218e-05,
+      "loss": 0.0,
+      "step": 4200
+    },
+    {
+      "epoch": 1.6673205178501374,
+      "grad_norm": 0.0,
+      "learning_rate": 1.6653589642997255e-05,
+      "loss": 0.0,
+      "step": 4250
+    },
+    {
+      "epoch": 1.6869360533542566,
+      "grad_norm": 0.0,
+      "learning_rate": 1.5672812867791292e-05,
+      "loss": 0.0,
+      "step": 4300
+    },
+    {
+      "epoch": 1.706551588858376,
+      "grad_norm": 0.0,
+      "learning_rate": 1.4692036092585329e-05,
+      "loss": 0.0,
+      "step": 4350
+    },
+    {
+      "epoch": 1.7261671243624952,
+      "grad_norm": 0.0,
+      "learning_rate": 1.3711259317379366e-05,
+      "loss": 0.0,
+      "step": 4400
+    },
+    {
+      "epoch": 1.7457826598666144,
+      "grad_norm": 0.0,
+      "learning_rate": 1.2730482542173403e-05,
+      "loss": 0.0,
+      "step": 4450
+    },
+    {
+      "epoch": 1.7653981953707336,
+      "grad_norm": 0.0,
+      "learning_rate": 1.174970576696744e-05,
+      "loss": 0.0,
+      "step": 4500
+    },
+    {
+      "epoch": 1.7850137308748528,
+      "grad_norm": 0.0,
+      "learning_rate": 1.0768928991761477e-05,
+      "loss": 0.0,
+      "step": 4550
+    },
+    {
+      "epoch": 1.804629266378972,
+      "grad_norm": 0.0,
+      "learning_rate": 9.788152216555512e-06,
+      "loss": 0.0,
+      "step": 4600
+    },
+    {
+      "epoch": 1.8242448018830915,
+      "grad_norm": 0.0,
+      "learning_rate": 8.80737544134955e-06,
+      "loss": 0.0,
+      "step": 4650
+    },
+    {
+      "epoch": 1.8438603373872107,
+      "grad_norm": 0.0,
+      "learning_rate": 7.826598666143586e-06,
+      "loss": 0.0,
+      "step": 4700
+    },
+    {
+      "epoch": 1.86347587289133,
+      "grad_norm": 0.0,
+      "learning_rate": 6.845821890937623e-06,
+      "loss": 0.0,
+      "step": 4750
+    },
+    {
+      "epoch": 1.8830914083954493,
+      "grad_norm": 0.0,
+      "learning_rate": 5.86504511573166e-06,
+      "loss": 0.0,
+      "step": 4800
+    },
+    {
+      "epoch": 1.9027069438995685,
+      "grad_norm": 0.0,
+      "learning_rate": 4.8842683405256965e-06,
+      "loss": 0.0,
+      "step": 4850
+    },
+    {
+      "epoch": 1.9223224794036877,
+      "grad_norm": 0.0,
+      "learning_rate": 3.9034915653197335e-06,
+      "loss": 0.0,
+      "step": 4900
+    },
+    {
+      "epoch": 1.941938014907807,
+      "grad_norm": 0.0,
+      "learning_rate": 2.92271479011377e-06,
+      "loss": 0.0,
+      "step": 4950
+    },
+    {
+      "epoch": 1.9615535504119261,
+      "grad_norm": 0.0,
+      "learning_rate": 1.941938014907807e-06,
+      "loss": 0.0,
+      "step": 5000
+    },
+    {
+      "epoch": 1.9811690859160453,
+      "grad_norm": 0.0,
+      "learning_rate": 9.61161239701844e-07,
+      "loss": 0.0,
+      "step": 5050
+    },
+    {
+      "epoch": 2.0,
+      "eval_loss": 0.0,
+      "eval_runtime": 4.9822,
+      "eval_samples_per_second": 33.72,
+      "eval_steps_per_second": 1.204,
+      "step": 5098
+    }
+  ],
+  "logging_steps": 50,
+  "max_steps": 5098,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 2,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.0,
+  "train_batch_size": 16,
+  "trial_name": null,
+  "trial_params": null
+}

nvidia-llama-nemotron-rerank-1b-v2_llama-E2-Pairwise/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5faa7d5c61b4800ee7490e3c72bda98672158b8917ee8959ced0ed37f02f6b26
+size 5265