Upload SPLADE-PT-BR model v1.0.0

README.md

@@ -77,14 +77,16 @@ pip install torch transformers
 
 ### Basic Usage
 
+**Option 1: Using HuggingFace Hub (Recommended)**
+
 ```python
 import torch
 from transformers import AutoTokenizer
-from splade.models.transformer_rep import Splade
+from modeling_splade import Splade
 
 # Load model and tokenizer
 model = Splade.from_pretrained("AxelPCG/splade-pt-br")
-tokenizer = AutoTokenizer.from_pretrained("neuralmind/bert-base-portuguese-cased")
+tokenizer = AutoTokenizer.from_pretrained("AxelPCG/splade-pt-br")
 model.eval()
 
 # Encode a query
@@ -109,6 +111,17 @@ values = query_vec[indices].tolist()
 print(f"Active dimensions: {len(indices)} / {query_vec.shape[0]}")
 ```
 
+**Option 2: Using SPLADE Library**
+
+```python
+from splade.models.transformer_rep import Splade
+from transformers import AutoTokenizer
+
+# Load model by pointing to HuggingFace repo
+model = Splade(model_type_or_dir="AxelPCG/splade-pt-br", agg="max", fp16=True)
+tokenizer = AutoTokenizer.from_pretrained("AxelPCG/splade-pt-br")
+```
+
 ## Limitations and Bias
 
 - Model trained on machine-translated Portuguese data (mMARCO)

config.json

@@ -4,6 +4,7 @@
   ],
   "model_type": "splade",
   "base_model": "neuralmind/bert-base-portuguese-cased",
+  "model_type_or_dir": "neuralmind/bert-base-portuguese-cased",
   "vocab_size": 29794,
   "hidden_size": 768,
   "num_hidden_layers": 12,
@@ -16,6 +17,12 @@
   "type_vocab_size": 2,
   "initializer_range": 0.02,
   "layer_norm_eps": 1e-12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "use_cache": true,
+  "classifier_dropout": null,
   "aggregation": "max",
-  "fp16": true
+  "output": "MLM",
+  "fp16": true,
+  "agg": "max"
 }

model_metadata.json

@@ -4,19 +4,17 @@
   "description": "SPLADE sparse retrieval model trained for Brazilian Portuguese",
   "author": "AxelPCG",
   "release_date": "2025-12-01",
-  
   "base_model": {
     "name": "neuralmind/bert-base-portuguese-cased",
     "type": "BERTimbau",
     "language": "Portuguese (Brazilian)",
     "vocab_size": 29794
   },
-  
   "training": {
     "training_dataset": "mMARCO Portuguese (unicamp-dl/mmarco)",
     "validation_dataset": "mRobust (unicamp-dl/mrobust)",
     "num_iterations": 150000,
-    "final_loss": 0.000047,
+    "final_loss": 4.7e-05,
     "batch_size": 8,
     "effective_batch_size": 32,
     "gradient_accumulation_steps": 4,
@@ -27,7 +25,6 @@
     "fp16": true,
     "optimizer": "AdamW",
     "scheduler": "linear_with_warmup",
-    
     "regularization": {
       "type": "FLOPS",
       "lambda_q": 0.0003,
@@ -35,7 +32,6 @@
       "T": 50000
     }
   },
-  
   "model_specs": {
     "architecture": "SPLADE",
     "aggregation": "max",
@@ -44,16 +40,19 @@
     "avg_active_dims_query": 120,
     "avg_active_dims_doc": 150
   },
-  
   "performance": {
-    "note": "Metrics will be updated after complete evaluation",
-    "expected": {
-      "MRR@10": "0.25-0.35",
-      "Recall@100": "0.85-0.95",
-      "Recall@1000": "0.95-0.99"
+    "dataset": "mRobust (TREC Robust04 Portuguese)",
+    "num_documents": 528032,
+    "num_queries": 250,
+    "metrics": {
+      "MRR@10": 0.453,
+      "evaluation_date": "2025-12-02"
+    },
+    "comparison": {
+      "splade_en_mrr10": 0.383,
+      "improvement": "+18.3%"
     }
   },
-  
   "usage": {
     "primary_use_case": "Sparse vector retrieval for Portuguese RAG systems",
     "recommended_for": [
@@ -68,21 +67,18 @@
       "custom": "Standard inverted index on non-zero dimensions"
     }
   },
-  
   "files": {
     "checkpoint": "model_final_checkpoint.tar",
     "config": "config.yaml",
     "tokenizer": "neuralmind/bert-base-portuguese-cased",
     "size_mb": 450
   },
-  
   "huggingface": {
     "repo_id": "AxelPCG/splade-pt-br",
     "model_type": "splade",
     "pipeline_tag": "feature-extraction",
     "license": "apache-2.0"
   },
-  
   "comparison_with_original": {
     "original_model": "SPLADE++",
     "original_language": "English",
@@ -94,16 +90,13 @@
       "Better semantic understanding of Brazilian Portuguese"
     ]
   },
-  
   "limitations": [
     "Optimized for Brazilian Portuguese",
     "Not tested on European Portuguese",
     "May require domain adaptation for specialized fields",
     "Max sequence length: 256 tokens"
   ],
-  
   "citation": {
     "bibtex": "@misc{splade-pt-br-2025, author = {Axel Chepanski}, title = {SPLADE-PT-BR: Sparse Retrieval for Portuguese}, year = {2025}, publisher = {Hugging Face}, url = {https://huggingface.co/AxelPCG/splade-pt-br}}"
   }
-}
-
+}

modeling_splade.py

@@ -0,0 +1,105 @@
+"""
+SPLADE Model for HuggingFace Hub
+Adapted from: https://github.com/naver/splade
+"""
+
+import torch
+from transformers import AutoModelForMaskedLM, PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutput
+
+
+class SpladeConfig(PretrainedConfig):
+    """Configuration class for SPLADE model"""
+    model_type = "splade"
+    
+    def __init__(
+        self,
+        base_model="neuralmind/bert-base-portuguese-cased",
+        aggregation="max",
+        fp16=True,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.base_model = base_model
+        self.aggregation = aggregation
+        self.fp16 = fp16
+
+
+class Splade(PreTrainedModel):
+    """
+    SPLADE model for sparse retrieval.
+    
+    This model produces sparse representations by:
+    1. Using a MLM head to get vocabulary-sized logits
+    2. Applying log(1 + ReLU(logits))
+    3. Max-pooling over sequence length
+    
+    Usage:
+        from transformers import AutoTokenizer
+        from modeling_splade import Splade
+        
+        model = Splade.from_pretrained("AxelPCG/splade-pt-br")
+        tokenizer = AutoTokenizer.from_pretrained("AxelPCG/splade-pt-br")
+        
+        # Encode query
+        query_tokens = tokenizer("Qual é a capital do Brasil?", return_tensors="pt")
+        with torch.no_grad():
+            query_vec = model(q_kwargs=query_tokens)["q_rep"]
+    """
+    config_class = SpladeConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        
+        # Load base BERT model with MLM head
+        base_model = getattr(config, 'base_model', 'neuralmind/bert-base-portuguese-cased')
+        self.transformer = AutoModelForMaskedLM.from_pretrained(base_model)
+        self.aggregation = getattr(config, 'aggregation', 'max')
+        self.fp16 = getattr(config, 'fp16', True)
+        
+    def encode(self, tokens):
+        """Encode tokens to sparse representation"""
+        # Get MLM logits
+        out = self.transformer(**tokens)
+        logits = out.logits  # shape (bs, seq_len, vocab_size)
+        
+        # Apply log(1 + ReLU(x))
+        relu_log = torch.log1p(torch.relu(logits))
+        
+        # Apply attention mask
+        attention_mask = tokens["attention_mask"].unsqueeze(-1)
+        masked = relu_log * attention_mask
+        
+        # Aggregate (max or sum)
+        if self.aggregation == "max":
+            values, _ = torch.max(masked, dim=1)
+            return values
+        else:  # sum
+            return torch.sum(masked, dim=1)
+    
+    def forward(self, q_kwargs=None, d_kwargs=None, **kwargs):
+        """
+        Forward pass supporting both query and document encoding.
+        
+        Args:
+            q_kwargs: Query tokens (dict with input_ids, attention_mask)
+            d_kwargs: Document tokens (dict with input_ids, attention_mask)
+            **kwargs: Additional arguments (for compatibility)
+        
+        Returns:
+            dict with 'q_rep' and/or 'd_rep' keys containing sparse vectors
+        """
+        output = {}
+        
+        if q_kwargs is not None:
+            output["q_rep"] = self.encode(q_kwargs)
+        
+        if d_kwargs is not None:
+            output["d_rep"] = self.encode(d_kwargs)
+        
+        # If neither q_kwargs nor d_kwargs, use kwargs directly
+        if not output and kwargs:
+            output["rep"] = self.encode(kwargs)
+        
+        return output

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fc862991df523373e5698b3341dd0a99245cd3590a345c2173170bc44b7cb6f0
-size 1307742766
+oid sha256:9fbf5a3a20f6e1d1ae82d2189d6f754434fca7e1b351ad76bbc3a90e6ab32587
+size 435884875

special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

tokenizer_config.json

@@ -1,6 +1 @@
-{
-  "tokenizer_class": "BertTokenizer",
-  "do_lower_case": false,
-  "model_max_length": 256,
-  "tokenizer_type": "neuralmind/bert-base-portuguese-cased"
-}
+{"do_lower_case": false, "init_inputs": []}