Upload folder using huggingface_hub

config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "RexRerankerModel"
+  ],
+  "backbone_name": "thebajajra/RexBERT-base",
+  "dropout": 0.0,
+  "dtype": "bfloat16",
+  "hidden_size": 768,
+  "model_type": "rex_reranker",
+  "num_bins": 11,
+  "pooling_strategy": "mean",
+  "sigma_delta": 0.08,
+  "sigma_max": 0.12,
+  "sigma_min": 0.04,
+  "transformers_version": "4.57.3",
+  "transitions": [
+    0.2,
+    0.5,
+    0.8
+  ],
+  "num_labels": 1,
+  "torch_dtype": "bfloat16",
+  "auto_map": {
+    "AutoConfig": "modeling_rex_reranker.RexRerankerConfig",
+    "AutoModel": "modeling_rex_reranker.RexRerankerModel",
+    "AutoModelForSequenceClassification": "modeling_rex_reranker.RexRerankerModel"
+  }
+}

model.safetensors

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

modeling_rex_reranker.py

@@ -0,0 +1,227 @@
+"""RexReranker Model for HuggingFace.
+
+Compatible with:
+- Transformers: AutoModel.from_pretrained(..., trust_remote_code=True)
+- Sentence Transformers: CrossEncoder(..., trust_remote_code=True)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, List, Union
+from dataclasses import dataclass
+
+from transformers import PretrainedConfig, PreTrainedModel, AutoModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+
+@dataclass
+class RexRerankerOutput(SequenceClassifierOutput):
+    """Output class for RexReranker with additional distributional information."""
+    loss: Optional[torch.Tensor] = None
+    logits: torch.Tensor = None  # Single relevance score [B, 1] for CrossEncoder compatibility
+    distribution_logits: torch.Tensor = None  # Full distribution [B, num_bins]
+    relevance: torch.Tensor = None  # Convenience: same as logits.squeeze(-1)
+    variance: torch.Tensor = None  # Prediction variance
+    entropy: torch.Tensor = None  # Distribution entropy
+
+
+class RexRerankerConfig(PretrainedConfig):
+    """Configuration for RexReranker model."""
+    
+    model_type = "rex_reranker"
+    
+    def __init__(
+        self,
+        backbone_name: str = "thebajajra/RexBERT-mini",
+        num_bins: int = 11,
+        dropout: float = 0.0,
+        pooling_strategy: str = "mean",
+        hidden_size: int = None,
+        num_labels: int = 1,  # CrossEncoder compatibility
+        transitions: List[float] = None,
+        sigma_min: float = 0.04,
+        sigma_max: float = 0.12,
+        sigma_delta: float = 0.08,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.backbone_name = backbone_name
+        self.num_bins = num_bins
+        self.dropout = dropout
+        self.pooling_strategy = pooling_strategy
+        self.hidden_size = hidden_size
+        self.num_labels = num_labels
+        self.transitions = transitions or [0.2, 0.5, 0.8]
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma_delta = sigma_delta
+
+
+class RexRerankerModel(PreTrainedModel):
+    """
+    RexBERT-based distributional reranker.
+    
+    Predicts a categorical distribution over K bins in [0, 1] representing
+    relevance scores. The output logits contain a single relevance score
+    for CrossEncoder compatibility, while the full distribution is available
+    via distribution_logits or predict_with_uncertainty().
+    
+    Compatible with:
+    - sentence_transformers.CrossEncoder
+    - transformers.AutoModelForSequenceClassification
+    """
+    
+    config_class = RexRerankerConfig
+    base_model_prefix = "rex_reranker"
+    supports_gradient_checkpointing = True
+    
+    def __init__(self, config: RexRerankerConfig):
+        super().__init__(config)
+        
+        assert config.pooling_strategy in ("cls", "mean")
+        self.pooling_strategy = config.pooling_strategy
+        self.num_bins = config.num_bins
+        
+        self.backbone = AutoModel.from_pretrained(
+            config.backbone_name,
+            trust_remote_code=True,
+        )
+        
+        if hasattr(self.backbone, "config") and hasattr(self.backbone.config, "use_cache"):
+            self.backbone.config.use_cache = False
+        
+        hidden_size = config.hidden_size or getattr(self.backbone.config, "hidden_size", None)
+        if hidden_size is None:
+            raise ValueError("Could not infer hidden_size.")
+        
+        self.dropout = nn.Dropout(config.dropout)
+        self.score_head = nn.Linear(hidden_size, config.num_bins)
+        
+        self.register_buffer(
+            "bin_centers",
+            torch.linspace(0.0, 1.0, config.num_bins),
+            persistent=False,
+        )
+        
+        self.post_init()
+    
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        output_distribution: bool = False,
+        **kwargs,  # Accept extra kwargs for CrossEncoder compatibility
+    ) -> Union[RexRerankerOutput, tuple]:
+        """
+        Forward pass.
+        
+        Args:
+            input_ids: Token IDs [B, T]
+            attention_mask: Attention mask [B, T]
+            labels: Optional relevance labels [B]
+            return_dict: Whether to return a dataclass
+            output_distribution: If True, include full distribution info in output
+            
+        Returns:
+            RexRerankerOutput with:
+                - logits: [B, 1] single relevance score (CrossEncoder compatible)
+                - distribution_logits: [B, num_bins] full distribution (if output_distribution=True)
+                - relevance, variance, entropy: convenience fields (if output_distribution=True)
+        """
+        out = self.backbone(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=True,
+        )
+        last_hidden = out.last_hidden_state
+        
+        if self.pooling_strategy == "cls":
+            pooled = last_hidden[:, 0, :]
+        else:
+            mask = attention_mask.unsqueeze(-1).float()
+            summed = (last_hidden * mask).sum(dim=1)
+            lengths = mask.sum(dim=1).clamp(min=1e-9)
+            pooled = summed / lengths
+        
+        # Get distribution logits
+        dist_logits = self.score_head(self.dropout(pooled))  # [B, num_bins]
+        
+        # Convert to single relevance score (expected value)
+        probs = F.softmax(dist_logits, dim=-1)
+        relevance = (probs * self.bin_centers.view(1, -1)).sum(dim=-1)  # [B]
+        
+        # Output single score as logits for CrossEncoder compatibility [B, 1]
+        logits = relevance.unsqueeze(-1)
+        
+        loss = None
+        if labels is not None:
+            loss = F.mse_loss(relevance, labels.float())
+        
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+        
+        # Compute additional stats if requested
+        variance = None
+        entropy = None
+        if output_distribution:
+            variance = (probs * (self.bin_centers.view(1, -1) - relevance.unsqueeze(-1)) ** 2).sum(dim=-1)
+            entropy = -(probs * torch.log(probs.clamp(min=1e-9))).sum(dim=-1)
+        
+        return RexRerankerOutput(
+            loss=loss,
+            logits=logits,
+            distribution_logits=dist_logits if output_distribution else None,
+            relevance=relevance,
+            variance=variance,
+            entropy=entropy,
+        )
+    
+    def predict_relevance(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Get relevance scores directly. Returns [B] tensor."""
+        outputs = self.forward(input_ids=input_ids, attention_mask=attention_mask)
+        return outputs.relevance
+    
+    def predict_with_uncertainty(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> dict:
+        """
+        Get relevance prediction with full uncertainty estimates.
+        
+        Returns:
+            dict with:
+                - relevance: [B] predicted relevance scores
+                - variance: [B] prediction variance (higher = more uncertain)
+                - entropy: [B] distribution entropy (higher = more uncertain)
+                - probs: [B, num_bins] full probability distribution
+                - distribution_logits: [B, num_bins] raw logits
+        """
+        outputs = self.forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_distribution=True,
+        )
+        probs = F.softmax(outputs.distribution_logits, dim=-1)
+        
+        return {
+            "relevance": outputs.relevance,
+            "variance": outputs.variance,
+            "entropy": outputs.entropy,
+            "probs": probs,
+            "distribution_logits": outputs.distribution_logits,
+        }

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,945 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "1": {
+      "content": "<|padding|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50254": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50255": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50256": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50257": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50280": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50281": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50282": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50283": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50284": {
+      "content": "[MASK]",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50285": {
+      "content": "[unused0]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50286": {
+      "content": "[unused1]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50287": {
+      "content": "[unused2]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50288": {
+      "content": "[unused3]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50289": {
+      "content": "[unused4]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50290": {
+      "content": "[unused5]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50291": {
+      "content": "[unused6]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50292": {
+      "content": "[unused7]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50293": {
+      "content": "[unused8]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50294": {
+      "content": "[unused9]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50295": {
+      "content": "[unused10]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50296": {
+      "content": "[unused11]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50297": {
+      "content": "[unused12]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50298": {
+      "content": "[unused13]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50299": {
+      "content": "[unused14]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50300": {
+      "content": "[unused15]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50301": {
+      "content": "[unused16]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50302": {
+      "content": "[unused17]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50303": {
+      "content": "[unused18]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50304": {
+      "content": "[unused19]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50305": {
+      "content": "[unused20]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50306": {
+      "content": "[unused21]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50307": {
+      "content": "[unused22]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50308": {
+      "content": "[unused23]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50309": {
+      "content": "[unused24]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50310": {
+      "content": "[unused25]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50311": {
+      "content": "[unused26]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50312": {
+      "content": "[unused27]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50313": {
+      "content": "[unused28]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50314": {
+      "content": "[unused29]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50315": {
+      "content": "[unused30]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50316": {
+      "content": "[unused31]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50317": {
+      "content": "[unused32]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50318": {
+      "content": "[unused33]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50319": {
+      "content": "[unused34]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50320": {
+      "content": "[unused35]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50321": {
+      "content": "[unused36]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50322": {
+      "content": "[unused37]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50323": {
+      "content": "[unused38]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50324": {
+      "content": "[unused39]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50325": {
+      "content": "[unused40]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50326": {
+      "content": "[unused41]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50327": {
+      "content": "[unused42]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50328": {
+      "content": "[unused43]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50329": {
+      "content": "[unused44]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50330": {
+      "content": "[unused45]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50331": {
+      "content": "[unused46]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50332": {
+      "content": "[unused47]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50333": {
+      "content": "[unused48]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50334": {
+      "content": "[unused49]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50335": {
+      "content": "[unused50]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50336": {
+      "content": "[unused51]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50337": {
+      "content": "[unused52]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50338": {
+      "content": "[unused53]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50339": {
+      "content": "[unused54]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50340": {
+      "content": "[unused55]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50341": {
+      "content": "[unused56]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50342": {
+      "content": "[unused57]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50343": {
+      "content": "[unused58]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50344": {
+      "content": "[unused59]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50345": {
+      "content": "[unused60]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50346": {
+      "content": "[unused61]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50347": {
+      "content": "[unused62]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50348": {
+      "content": "[unused63]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50349": {
+      "content": "[unused64]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50350": {
+      "content": "[unused65]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50351": {
+      "content": "[unused66]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50352": {
+      "content": "[unused67]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50353": {
+      "content": "[unused68]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50354": {
+      "content": "[unused69]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50355": {
+      "content": "[unused70]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50356": {
+      "content": "[unused71]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50357": {
+      "content": "[unused72]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50358": {
+      "content": "[unused73]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50359": {
+      "content": "[unused74]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50360": {
+      "content": "[unused75]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50361": {
+      "content": "[unused76]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50362": {
+      "content": "[unused77]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50363": {
+      "content": "[unused78]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50364": {
+      "content": "[unused79]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50365": {
+      "content": "[unused80]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50366": {
+      "content": "[unused81]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50367": {
+      "content": "[unused82]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 7999,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

utils.py

@@ -0,0 +1,152 @@
+"""
+RexReranker Inference Utilities.
+
+This module provides helper functions for converting model logits to relevance scores.
+The model outputs logits for 11 bins representing a distribution over [0, 1].
+To get a relevance score, apply softmax and compute the expected value.
+
+Example usage:
+    from transformers import AutoModelForSequenceClassification, AutoTokenizer
+    from utils import logits_to_relevance, logits_to_relevance_with_uncertainty
+    import torch
+    
+    model = AutoModelForSequenceClassification.from_pretrained("path/to/model")
+    tokenizer = AutoTokenizer.from_pretrained("path/to/model")
+    
+    inputs = tokenizer(
+        "Query: best laptop",
+        "Title: MacBook Pro\nDescription: Great laptop for developers",
+        return_tensors="pt",
+        truncation=True,
+    )
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        
+        # Simple relevance score
+        relevance = logits_to_relevance(outputs.logits)
+        print(f"Relevance: {relevance.item():.3f}")
+        
+        # With uncertainty estimates
+        result = logits_to_relevance_with_uncertainty(outputs.logits)
+        print(f"Relevance: {result['relevance'].item():.3f}")
+        print(f"Variance: {result['variance'].item():.4f}")
+        print(f"Entropy: {result['entropy'].item():.3f}")
+"""
+
+import torch
+from typing import Dict
+
+
+# Configuration
+NUM_BINS = 11
+BIN_CENTERS = torch.linspace(0.0, 1.0, NUM_BINS)
+
+
+def logits_to_relevance(logits: torch.Tensor) -> torch.Tensor:
+    """
+    Convert model logits to relevance scores.
+    
+    Args:
+        logits: Model output logits [B, 11]
+        
+    Returns:
+        relevance: Relevance scores [B] in range [0, 1]
+    """
+    probs = torch.softmax(logits, dim=-1)
+    bin_centers = BIN_CENTERS.to(logits.device)
+    return (probs * bin_centers.view(1, -1)).sum(dim=-1)
+
+
+def logits_to_relevance_with_uncertainty(logits: torch.Tensor) -> Dict[str, torch.Tensor]:
+    """
+    Convert model logits to relevance scores with uncertainty estimates.
+    
+    Args:
+        logits: Model output logits [B, 11]
+        
+    Returns:
+        dict with:
+            - relevance: [B] predicted relevance scores in [0, 1]
+            - variance: [B] prediction variance (higher = more uncertain)
+            - entropy: [B] distribution entropy (higher = more uncertain)
+            - probs: [B, 11] full probability distribution over bins
+    """
+    probs = torch.softmax(logits, dim=-1)
+    bin_centers = BIN_CENTERS.to(logits.device)
+    
+    relevance = (probs * bin_centers.view(1, -1)).sum(dim=-1)
+    variance = (probs * (bin_centers.view(1, -1) - relevance.unsqueeze(-1)) ** 2).sum(dim=-1)
+    entropy = -(probs * torch.log(probs.clamp(min=1e-9))).sum(dim=-1)
+    
+    return {
+        "relevance": relevance,
+        "variance": variance,
+        "entropy": entropy,
+        "probs": probs,
+    }
+
+
+def batch_rerank(
+    model,
+    tokenizer,
+    query: str,
+    documents: list,
+    max_length: int = 2048,
+    batch_size: int = 32,
+    device: str = None,
+) -> list:
+    """
+    Rerank a list of documents for a given query.
+    
+    Args:
+        model: The RexReranker model
+        tokenizer: The tokenizer
+        query: The search query
+        documents: List of dicts with 'title' and 'description' keys
+        max_length: Maximum sequence length
+        batch_size: Batch size for inference
+        device: Device to use (default: auto-detect)
+        
+    Returns:
+        List of dicts with original document info plus 'relevance', 'variance', 'entropy'
+    """
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    model = model.to(device)
+    model.eval()
+    
+    results = []
+    
+    for i in range(0, len(documents), batch_size):
+        batch_docs = documents[i:i + batch_size]
+        
+        # Format inputs
+        texts_a = [f"Query: {query}" for _ in batch_docs]
+        texts_b = [f"Title: {doc.get('title', '')}\nDescription: {doc.get('description', '')}" for doc in batch_docs]
+        
+        inputs = tokenizer(
+            texts_a,
+            texts_b,
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors="pt",
+        ).to(device)
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+            batch_results = logits_to_relevance_with_uncertainty(outputs.logits)
+        
+        for j, doc in enumerate(batch_docs):
+            results.append({
+                **doc,
+                "relevance": batch_results["relevance"][j].item(),
+                "variance": batch_results["variance"][j].item(),
+                "entropy": batch_results["entropy"][j].item(),
+            })
+    
+    # Sort by relevance (descending)
+    results.sort(key=lambda x: x["relevance"], reverse=True)
+    return results