Update README.md

README.md

@@ -14,118 +14,216 @@ tags:
 pipeline_tag: fill-mask
 library_name: transformers
 ---
-# [RexBERT-mini](https://huggingface.co/thebajajra/RexBERT-mini)
+# RexBERT-mini
 
-> An efficient, English encoder-only model (masked-language model) with ~8k token context, targeted at e-commerce and retail NLP.
+[![License: Apache2.0](https://img.shields.io/badge/License-Apache2.0-green.svg)](https://www.apache.org/licenses/LICENSE-2.0)
+[![Models](https://img.shields.io/badge/🤗%20Hugging%20Face-Models-red)](https://huggingface.co/collections/thebajajra/rexbert-68cc4b1b8a272f6beae5ebb8)
+[![Data](https://img.shields.io/badge/🤗%20Training%20Data-Ecomniverse-yellow)](https://huggingface.co/datasets/thebajajra/Ecom-niverse)
+[![GitHub](https://img.shields.io/badge/GitHub-Code-blue)](https://github.com/bajajra/RexBERT)
+
+> **TL;DR**: An encoder-only transformer (ModernBERT-style) for **e-commerce** applications, trained in three phases—**Pre-training**, **Context Extension**, and **Decay**—to power product search, attribute extraction, classification, and embeddings use cases. The model has been trained on 2.3T+ tokens along with 350B+ e-commerce-specific tokens
 
 ---
 
-## Model summary
+## Table of Contents
+- [Quick Start](#quick-start)
+- [Intended Uses & Limitations](#intended-uses--limitations)
+- [Model Description](#model-description)
+- [Training Recipe](#training-recipe)
+- [Data Overview](#data-overview)
+- [Evaluation](#evaluation)
+- [Usage Examples](#usage-examples)
+  - [Masked language modeling](#1-masked-language-modeling)
+  - [Embeddings / feature extraction](#2-embeddings--feature-extraction)
+  - [Text classification fine-tune](#3-text-classification-fine-tune)
+- [Model Architecture & Compatibility](#model-architecture--compatibility)
+- [Efficiency & Deployment Tips](#efficiency--deployment-tips)
+- [Responsible & Safe Use](#responsible--safe-use)
+- [License](#license)
+- [Maintainers & Contact](#maintainers--contact)
+- [Citation](#citation)
 
-- **Model type:** `ModernBertForMaskedLM` (encoder-only, masked-language modeling head)  
-- **Domain:** **e-commerce**/**retail**/**shopping**
-- **Language:** English 
-- **Context length:** 7,999–8,192 tokens (config max_position_embeddings=7999; ModernBERT supports up to 8192)  
-- **License:** Apache-2.0  
 ---
 
-## Intended uses & limitations
+## Quick Start
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM, pipeline
+
+MODEL_ID = "thebajajra/RexBERT-mini"
 
-### Direct use
-- **Fill-mask** and cloze completion (e.g., product titles, attributes, query reformulation).  
-- **Embeddings / feature extraction** for classification, clustering, retrieval re-ranking, and semantic search in retail catalogs and queries (via pooled encoder states). (ModernBERT is a drop-in BERT-style encoder.)
+# Tokenizer
+tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
 
-### Downstream use
-- Fine-tune for product categorization, attribute extraction, NER, intent classification, and retrieval-augmented ranking tasks in commerce search & browse. (Use a task head or pooled embeddings.)
+# 1) Fill-Mask (if MLM head is present)
+mlm = pipeline("fill-mask", model=MODEL_ID, tokenizer=tok)
+print(mlm("These running shoes are great for [MASK] training."))
+
+# 2) Feature extraction (CLS or mean-pooled embeddings)
+enc = AutoModel.from_pretrained(MODEL_ID)
+inputs = tok(["wireless mouse", "ergonomic mouse pad"], padding=True, truncation=True, return_tensors="pt")
+with torch.no_grad():
+    out = enc(**inputs, output_hidden_states=True)
+# Mean-pool last hidden state for sentence embeddings
+emb = (out.last_hidden_state * inputs.attention_mask.unsqueeze(-1)).sum(dim=1) / inputs.attention_mask.sum(dim=1, keepdim=True)
+```
 
-### Out-of-scope / not recommended
-- **Autoregressive text generation** or chat; this is not a decoder LLM. Use decoder or seq2seq models for long-form generation.
 
 ---
 
-## How to get started
+## Intended Uses & Limitations
 
-```python
-from transformers import AutoTokenizer, AutoModelForMaskedLM
+**Use cases**
+- Product & query **retrieval/semantic search** (titles, descriptions, attributes)
+- **Attribute extraction** / slot filling (brand, color, size, material)
+- **Classification** (category assignment, unsafe/regulated item filtering, review sentiment)
+- **Reranking** and **query understanding** (spelling/ASR normalization, acronym expansion)
 
-model_id = "thebajajra/RexBERT-mini"
-tok = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForMaskedLM.from_pretrained(model_id)
+**Out of scope**
+- Long-form **generation** (use a decoder/seq-to-seq LM instead)
+- High-stakes decisions without human review (pricing, compliance, safety flags)
 
-text = "The customer purchased a [MASK] with free shipping."
-inputs = tok(text, return_tensors="pt")
-logits = model(**inputs).logits  # use top-k on tok.mask_token_id
-```
+**Target users**
+- Search/recs engineers, e-commerce data teams, ML researchers working on domain-specific encoders
 
 ---
 
-## Model details
+## Model Description
 
-### Architecture (from config)
-- **Backbone:** ModernBERT (`model_type: "modernbert"`, `architectures: ["ModernBertForMaskedLM"]`)  
-- **Layers / heads / width:** 19 encoder layers, 8 attention heads, hidden size 512; intermediate (MLP) size 768; GELU activations.  
-- **Attention:** Local window 128 with **global attention every 3 layers**; RoPE θ=160k (local & global).  
-- **Positional strategy:** `position_embedding_type: "sans_pos"`.  
-
-## Training data & procedure
+RexBERT-mini is an **encoder-only**, 68M parameter transformer trained with a masked-language-modeling objective and optimized for **e-commerce related text**. The three-phase training curriculum improves general language understanding, extends context handling, and then **specializes** on a very large corpus of commerce data to capture domain-specific terminology and entity distributions.
 
 ---
 
-## Performance Highlights
+## Training Recipe
+
+RexBERT-mini was trained in **three phases**:
 
-### MLM – Token Classification ([E-Commerce](https://github.com/amazon-science/esci-data))
+1) **Pre-training**  
+   General-purpose MLM pre-training on diverse English text for robust linguistic representations.
 
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/6893dd21467f7d2f5f358a95/BAafq5QMJI_-CQSgr5PzF.png)
+2) **Context Extension**  
+   Continued training with **increased max sequence length** to better handle long product pages, concatenated attribute blocks, multi-turn queries, and facet strings. This preserves prior capabilities while expanding context handling.
 
+3) **Decay on 350B+ e-commerce tokens**  
+   Final specialization stage on **350B+ domain-specific tokens** (product catalogs, queries, reviews, taxonomy/attributes). Learning rate and sampling weights are annealed (decayed) to consolidate domain knowledge and stabilize performance on commerce tasks.
 
-> RexBERT-mini outperforms DistilBERT on all token classification tasks by significant margin.
+**Training details (fill in):**
+- Optimizer / LR schedule: TODO
+- Effective batch size / steps per phase: TODO
+- Context lengths per phase (e.g., 512 → 1k/2k): TODO
+- Tokenizer/vocab: TODO
+- Hardware & wall-clock: TODO
+- Checkpoint tags: TODO (e.g., `pretrain`, `ext`, `decay`)
 
-#### Product Title
+---
 
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/6893dd21467f7d2f5f358a95/s3yMJ1KEGTfugCEivJTXD.png)
+## Data Overview
 
-#### Product Description
+- **Domain mix:**  
+- **Data quality:** 
 
 
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/6893dd21467f7d2f5f358a95/9sbtklAvcaNcprKqzVbty.png)
 
 ---
 
+## Evaluation
 
-## Technical notes for practitioners
 
-- **Pooling:** Use mean pooling over last hidden states (the config’s classifier pooling is `"mean"`), or task-specific pooling.  
-- **Long sequences:** Leverage the extended context for product pages, multi-turn queries, or concatenated fields; ModernBERT uses efficient attention and RoPE for long inputs.  
-- **Libraries:** Tested with `transformers>=4.48.0`
+### Performance Highlights 
 
----
 
-## Model sources
 
-- **Hugging Face:** `thebajajra/RexBERT-mini` — https://huggingface.co/thebajajra/RexBERT-mini  
-- **Background on ModernBERT:** https://huggingface.co/docs/transformers/en/model_doc/modernbert and overview: https://huggingface.co/docs/transformers/model_doc/modernbert
 
 ---
 
-## Citation
+## Usage Examples
 
-If you use this model, please cite the repository:
+### 1) Masked language modeling
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline
+
+m = AutoModelForMaskedLM.from_pretrained("thebajajra/RexBERT-mini")
+t = AutoTokenizer.from_pretrained("thebajajra/RexBERT-mini")
+fill = pipeline("fill-mask", model=m, tokenizer=t)
+
+fill("Best [MASK] headphones under $100.")
+```
 
+### 2) Embeddings / feature extraction
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tok = AutoTokenizer.from_pretrained("thebajajra/RexBERT-mini")
+enc = AutoModel.from_pretrained("thebajajra/RexBERT-mini")
+
+texts = ["nike air zoom pegasus 40", "running shoes pegasus zoom nike"]
+batch = tok(texts, padding=True, truncation=True, return_tensors="pt")
+
+with torch.no_grad():
+    out = enc(**batch)
+# Mean-pool last hidden state
+attn = batch["attention_mask"].unsqueeze(-1)
+emb = (out.last_hidden_state * attn).sum(1) / attn.sum(1)
+# Normalize for cosine similarity (recommended for retrieval)
+emb = torch.nn.functional.normalize(emb, p=2, dim=1)
 ```
-@software{rexbert_mini_2025,
-  title        = {RexBERT-mini},
-  author       = {},
-  year         = {2025},
-  url          = {}
-}
+
+### 3) Text classification fine-tune
+```python
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
+
+tok = AutoTokenizer.from_pretrained("thebajajra/RexBERT-mini")
+model = AutoModelForSequenceClassification.from_pretrained("thebajajra/RexBERT-mini", num_labels=NUM_LABELS)
+
+# Prepare your Dataset objects: train_ds, val_ds (text→label)
+args = TrainingArguments(
+    per_device_train_batch_size=32,
+    per_device_eval_batch_size=32,
+    learning_rate=3e-5,
+    num_train_epochs=3,
+    evaluation_strategy="steps",
+    fp16=True,
+    report_to="none",
+    load_best_model_at_end=True,
+)
+
+trainer = Trainer(model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds, tokenizer=tok)
+trainer.train()
 ```
 
 ---
 
-## Contact & maintenance
+## Model Architecture & Compatibility
+
+- **Architecture:** Encoder-only, ModernBERT-style **mini** model.  
+- **Libraries:** Works with **🤗 Transformers**; supports **fill-mask** and **feature-extraction** pipelines.  
+- **Context length:** Increased during the **Context Extension** phase—ensure `max_position_embeddings` in `config.json` matches your desired max length.  
+- **Files:** `config.json`, tokenizer files, and (optionally) heads for MLM or classification.  
+- **Export:** Standard PyTorch weights; you can export ONNX / TorchScript for production if needed.
+
+---
+
+## Responsible & Safe Use
+
+- **Biases:** Commerce data can encode brand, price, and region biases; audit downstream classifiers/retrievers for disparate error rates across categories/regions.
+- **Sensitive content:** Add filters for adult/regulated items; document moderation thresholds if you release classifiers.
+- **Privacy:** Do not expose PII; ensure training data complies with terms and applicable laws.
+- **Misuse:** This model is **not** a substitute for legal/compliance review for listings.
+
+---
+
+## License
+
+- **License:** `apache-2.0`.  
+---
+
+## Maintainers & Contact
 
-- **Author(s):** [Rahul Bajaj](https://huggingface.co/thebajajra)
+- **Author/maintainer:** [Rahul Bajaj](https://huggingface.co/thebajajra) 
+
+---
 
-- **Issues / questions:** Open an issue or discussion on the HF model page.
 
 ---