Upload folder using huggingface_hub

README.md

@@ -1,289 +1,53 @@
----
-license: apache-2.0
-language:
-- en
-tags:
-- ecommerce
-- e-commerce
-- retail
-- marketplace
-- shopping
-- amazon
-- ebay
-- alibaba
-- google
-- rakuten
-- bestbuy
-- walmart
-- flipkart
-- wayfair
-- shein
-- target
-- etsy
-- shopify
-- taobao
-- asos
-- carrefour
-- costco
-- overstock
-- pretraining
-- encoder
-- language-modeling
-- foundation-model
-base_model:
-- thebajajra/RexBERT-base
-pipeline_tag: text-ranking
-library_name: sentence-transformers
----
+# Reranker Model
 
-<p align="center">
-<img src="https://cdn-uploads.huggingface.co/production/uploads/6893dd21467f7d2f5f358a95/apOIbl5PdJuRk-tQMdDc8.png" alt="RexReranker">
-</p>
-<p align="center">
-</p>
+This model was exported from checkpoint: `rexbert-reranker-base/checkpoint-67628/`
 
-# RexReranker Base
+## Model Details
 
-A distributional **e-commerce** neural reranker based on RexBERT-base that predicts relevance scores as a probability distribution, providing both accurate relevance predictions and uncertainty estimates.
+- **Base Model**: thebajajra/RexBERT-base
+- **Task**: Document Reranking
+- **Output**: Relevance score between 0 and 1
 
-## Features
-
-- **Distributional Output**: Predicts a probability distribution over relevance bins (0.0 to 1.0)
-- **Uncertainty Estimates**: Provides variance and entropy for confidence assessment
-- **CrossEncoder Compatible**: Works directly with Sentence Transformers CrossEncoder
-- **Mean Pooling**: Uses mean pooling over all tokens for robust representations
-
-## Installation
-
-```bash
-pip install transformers sentence-transformers torch
-```
-
-## Quick Start
-
-### 1. Using HuggingFace Transformers
+## Usage
 
 ```python
-from transformers import AutoModel, AutoTokenizer
 import torch
+from transformers import AutoTokenizer
+from train_modernbert_reranker import ModernBERTReranker
 
 # Load model and tokenizer
-model = AutoModel.from_pretrained(
-    "thebajajra/RexReranker-base",
-    trust_remote_code=True
-)
-tokenizer = AutoTokenizer.from_pretrained("thebajajra/RexReranker-base")
-
-# Move to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
+model_path = "rexreranker-base"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = ModernBERTReranker.from_pretrained(model_path)
 model.eval()
 
-# Prepare input (query-document pair)
-query = "best laptop for programming"
-title = "MacBook Pro M3"
-description = "Powerful laptop with M3 chip, 16GB RAM, perfect for developers and creative professionals"
+# Example inference
+query = "wireless bluetooth headphones"
+document = "Title: Sony WH-1000XM5\nDescription: Premium wireless headphones with noise cancellation"
 
-inputs = tokenizer(
-    f"Query: {query}",
-    f"Title: {title}\nDescription: {description}",
-    return_tensors="pt",
-    truncation=True,
-    max_length=2048,
-).to(device)
+inputs = tokenizer(query, document, return_tensors="pt", truncation=True, max_length=2048)
 
-# Get relevance score
 with torch.no_grad():
-    score = model.predict_relevance(**inputs)
-    print(f"Relevance Score: {score.item():.4f}")
-```
-
-### 2. Using Sentence Transformers CrossEncoder
-
-```python
-from sentence_transformers import CrossEncoder
-
-# Load as CrossEncoder
-model = CrossEncoder(
-    "thebajajra/RexReranker-base",
-    trust_remote_code=True
-)
-
-# Single prediction
-query = "best laptop for programming"
-document = "MacBook Pro M3 - Powerful laptop with M3 chip for developers"
-
-score = model.predict([(query, document)])[0]
-print(f"Score: {score:.4f}")
-```
-
-### 3. Batch Reranking with CrossEncoder
-
-```python
-from sentence_transformers import CrossEncoder
-
-model = CrossEncoder("thebajajra/RexReranker-base", trust_remote_code=True)
-
-query = "best laptop for programming"
-documents = [
-    "MacBook Pro M3 - Powerful laptop with M3 chip for developers",
-    "Gaming Mouse RGB - High precision gaming mouse with 16000 DPI",
-    "ThinkPad X1 Carbon - Business ultrabook with long battery life",
-    "Mechanical Keyboard - Cherry MX switches for typing comfort",
-    "Dell XPS 15 - Premium laptop with 4K OLED display",
-]
-
-# Get scores for all documents
-pairs = [(query, doc) for doc in documents]
-scores = model.predict(pairs)
-
-# Print ranked results
-print(f"Query: {query}\n")
-for doc, score in sorted(zip(documents, scores), key=lambda x: x[1], reverse=True):
-    print(f"  {score:.4f} | {doc[:60]}")
-```
-
-### 4. Using CrossEncoder's rank() Method
-
-```python
-from sentence_transformers import CrossEncoder
-
-model = CrossEncoder("thebajajra/RexReranker-base", trust_remote_code=True)
-
-query = "wireless headphones with noise cancellation"
-documents = [
-    "Sony WH-1000XM5 - Industry-leading noise cancellation headphones",
-    "Apple AirPods Max - Premium over-ear headphones with spatial audio",
-    "Bose QuietComfort 45 - Comfortable wireless noise cancelling headphones",
-    "JBL Tune 750BTNC - Affordable wireless headphones with ANC",
-    "Logitech Gaming Headset - Wired gaming headphones with microphone",
-]
-
-# Rank documents
-results = model.rank(query, documents, top_k=3)
-
-print(f"Query: {query}\n")
-print("Top 3 Results:")
-for result in results:
-    idx = result['corpus_id']
-    score = result['score']
-    print(f"  {score:.4f} | {documents[idx][:60]}")
-```
-
-### 5. With Uncertainty Estimates
-
-```python
-from transformers import AutoModel, AutoTokenizer
-import torch
-
-model = AutoModel.from_pretrained("thebajajra/RexReranker-base", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("thebajajra/RexReranker-base")
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device).eval()
-
-# Prepare inputs
-inputs = tokenizer(
-    "Query: best laptop for programming",
-    "Title: MacBook Pro\nDescription: Great laptop for developers",
-    return_tensors="pt",
-    truncation=True,
-).to(device)
-
-# Get prediction with uncertainty
-with torch.no_grad():
-    result = model.predict_with_uncertainty(**inputs)
-
-print(f"Relevance: {result['relevance'].item():.4f}")
-print(f"Variance:  {result['variance'].item():.6f}")  # Higher = more uncertain
-print(f"Entropy:   {result['entropy'].item():.4f}")   # Higher = more uncertain
-
-# Access full probability distribution
-print(f"\nDistribution over bins:")
-probs = result['probs'][0].cpu().numpy()
-for i, p in enumerate(probs):
-    bin_center = i / (len(probs) - 1)
-    bar = "█" * int(p * 50)
-    print(f"  {bin_center:.1f}: {bar} ({p:.3f})")
-```
-
-### 6. Batch Processing for Production
-
-```python
-from transformers import AutoModel, AutoTokenizer
-import torch
-from torch.utils.data import DataLoader
-
-model = AutoModel.from_pretrained("thebajajra/RexReranker-base", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("thebajajra/RexReranker-base")
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device).eval()
-
-def rerank_batch(query: str, documents: list, batch_size: int = 32) -> list:
-    """Rerank documents for a query with batched inference."""
-    
-    # Prepare all inputs
-    all_inputs = []
-    for doc in documents:
-        title = doc.get("title", "")
-        description = doc.get("description", "")
-        inputs = tokenizer(
-            f"Query: {query}",
-            f"Title: {title}\nDescription: {description}",
-            truncation=True,
-            max_length=2048,
-            padding=False,
-        )
-        all_inputs.append(inputs)
-    
-    # Batch inference
-    all_scores = []
-    for i in range(0, len(all_inputs), batch_size):
-        batch = all_inputs[i:i + batch_size]
-        padded = tokenizer.pad(batch, return_tensors="pt").to(device)
-        
-        with torch.no_grad():
-            scores = model.predict_relevance(**padded)
-            all_scores.extend(scores.cpu().tolist())
+    outputs = model(**inputs)
+    score = outputs.logits.squeeze().item()
     
-    # Add scores to documents and sort
-    for doc, score in zip(documents, all_scores):
-        doc["score"] = score
-    
-    return sorted(documents, key=lambda x: x["score"], reverse=True)
-
-# Example usage
-query = "best laptop for programming"
-documents = [
-    {"title": "MacBook Pro M3", "description": "Powerful laptop for developers"},
-    {"title": "Gaming Mouse", "description": "High DPI gaming mouse"},
-    {"title": "ThinkPad X1", "description": "Business laptop with long battery"},
-]
-
-ranked = rerank_batch(query, documents)
-for doc in ranked:
-    print(f"{doc['score']:.4f} | {doc['title']}")
+print(f"Relevance score: {score:.4f}")
 ```
 
 ## Input Format
 
 The model expects query-document pairs formatted as:
+```
+Query: <query text>
+[SEP]
+Title: <title>
+Description: <description>
+```
 
-| Field | Format |
-|-------|--------|
-| Text A (Query) | `Query: {your search query}` |
-| Text B (Document) | `Title: {document title}\nDescription: {document description}` |
-
-## Output Details
-
-### Standard Output (CrossEncoder compatible)
-- `outputs.logits`: Shape `[B, 1]` - Single relevance score per example
-- `outputs.relevance`: Shape `[B]` - Same as logits squeezed
+## Training
 
-### With Uncertainty (`output_distribution=True` or `predict_with_uncertainty()`)
-- `relevance`: Expected relevance score [0, 1]
-- `variance`: Prediction variance (higher = less confident)
-- `entropy`: Distribution entropy (higher = less confident)
-- `probs`: Full probability distribution over bins
-- `distribution_logits`: Raw logits before softmax
+This model was trained on the Amazebay reranker dataset with:
+- Max sequence length: 2048
+- BF16 precision
+- Flash Attention 2
+- Multi-GPU training (4 GPUs)

config.json

@@ -1,28 +1,52 @@
 {
   "architectures": [
-    "RexRerankerModel"
+    "ModernBertForSequenceClassification"
   ],
-  "backbone_name": "thebajajra/RexBERT-base",
-  "dropout": 0.0,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 50281,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "mean",
+  "cls_token_id": 50281,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
   "dtype": "bfloat16",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 50282,
+  "global_attn_every_n_layers": 3,
+  "global_rope_theta": 160000.0,
+  "gradient_checkpointing": false,
+  "hidden_activation": "gelu",
   "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"
-  }
-}
+  "id2label": {
+    "0": "LABEL_0"
+  },
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 1152,
+  "label2id": {
+    "LABEL_0": 0
+  },
+  "layer_norm_eps": 1e-05,
+  "local_attention": 128,
+  "local_rope_theta": 160000.0,
+  "max_position_embeddings": 7999,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert",
+  "norm_bias": false,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 22,
+  "pad_token_id": 50283,
+  "position_embedding_type": "sans_pos",
+  "problem_type": "regression",
+  "repad_logits_with_grad": false,
+  "sep_token_id": 50282,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "transformers_version": "4.57.0",
+  "vocab_size": 50368
+}

model.safetensors

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

reranker_config.json

@@ -0,0 +1 @@
+{"use_regression": true, "model_type": "reranker"}

training_metadata.json

@@ -0,0 +1,8 @@
+{
+  "checkpoint_path": "rexbert-reranker-base/checkpoint-67628/",
+  "base_model": "thebajajra/RexBERT-base",
+  "global_step": 67628,
+  "epoch": 5.500447336315576,
+  "best_metric": 0.753167649044099,
+  "best_model_checkpoint": "./rexbert-reranker-base/checkpoint-64554"
+}