mlx-raclate 0.1.0b1__py3-none-any.whl

mlx_raclate/tuner/TUNER.md

@@ -0,0 +1,305 @@
+# RACLATE Tuner
+
+The `tuner` module is the training engine of RACLATE (**R**etrieval **A**nd **C**lassification including **LATE** interaction models). Built entirely on Apple's [MLX](https://github.com/ml-explore/mlx) framework, it provides a highly efficient, unified interface for fine-tuning *small* Transformer-based classifiers on Apple Silicon.
+
+This trainer supports standard dense retrieval, classification, and masked language modeling, as well as **Late Interaction (ColBERT-style)** training patterns.
+
+## Key Features
+
+*   **Apple Silicon Native:** Fully optimized for M-series chips using MLX.
+*   **Full Training:**  Full fine-tuning of pretrained models (_see supported architectures below_). LORA fine-tuning is not supported (yet). The library allows transfer learning, meaning that existing heads can be stripped out of pretrained models (and new heads can be added to base models for specific tasks)
+*   **Memory Efficiency:** Built-in support for **Gradient Accumulation** and **Gradient Checkpointing** to train larger batches/models on limited Unified Memory.
+*   **Flexible Schedulers:** Linear, Cosine, and Constant learning rate schedules with warmup.
+*   **Smart Collators:** Task-specific data collators that handle padding, masking, and chat templates automatically.
+*   **Embedding Freezing:** Option to freeze embedding layers to speed up fine-tuning or prevent catastrophic forgetting.
+*   **HF Hub Integration (TODO):** Seamless saving and pushing of checkpoints to the Hugging Face Hub.
+
+## Supported Architectures
+
+The trainer supports a variety of modern architectures supporting long context (relative to BERT models). As these models are meant to be trained and run on local machines, model implementations are specifically optimized for small-to-mid-sized models:
+
+*   **ModernBERT**: MLX implementation of `answerdotai/ModernBERT-base` (encoder-only). Long context (8k) and high efficiency.
+*   **Qwen 3**: MLX implementation of `Qwen/Qwen3-Embedding-0.6B` (32k context window) which leverages the qwen3 architecture.
+*   **Gemma 3**: MLX implementation of `google/embeddinggemma-300m` (2k context window) which leverages the gemma3 text variant architecture with a few tweaks. As per the official embeddinggemma3 architecture, the attention mask is set to causal or bi-directional based on a config parameter (`use_bidirectional_attn` or `use_bidirectional_attention`). Therefore, it is possible to switch between encoder and decoder mode, and standard gemma3_text models (32k context window) are also supported. 
+*   **T5Gemma-Encoder**: MLX implementation of `google/t5gemma-b-b-ul2`, but only keeping the encoder weights at initialization (the encoder config is merged into the main model config)
+*   **LFM2**: MLX implementation of `LiquidAI/LFM2-350M` (Causal/AR) which also supports `LiquidAI/LFM2-ColBERT-350M` when model config file includes `use_late_interaction=True`. These models have a context window of 128k tokens. In training mode, 128k tokens exceeds the RAM capacity of most Apple hardware. _See parameters below to cap sequences to a more reasonable length during training_
+
+
+## Supported Tasks & Pipelines
+
+The `Trainer` adapts its logic based on the `task_type` and the specific model class initialized.
+
+### 1. Sentence Similarity (Embedding & Retrieval)
+Train models for semantic search, clustering, or RAG.
+*   **Task Type:** `sentence-similarity` 
+*   **Training Modes:**
+    *   **Bi-Encoder (Dense):** Standard cosine similarity optimization.
+    *   **Late Interaction (MaxSim):** ColBERT-style interaction where fine-grained token-level similarities are computed (requires `use_late_interaction=True`).
+*   **Loss Functions:** Automatically selects between **MNRL (Multiple Negatives Ranking Loss)** for triplets/pairs or **MSE/Cosine Loss** for scored pairs.
+
+### 2. Sequence Classification
+Train discriminative models for sentiment analysis, intent detection, etc.
+*   **Task Type:** `text-classification`
+*   **Features:**
+    *   Supports Multi-class and Binary classification.
+    *   Supports Regression (if `is_regression=True`).
+    *   Native support for Chat Templates in tokenizer.
+
+### 3. Masked Language Modeling (MLM)
+Perform domain adaptation on raw text.
+*   **Task Type:** `masked-lm`
+*   **Features:** Implements the standard 80% mask / 10% random / 10% original masking strategy dynamically during training.
+
+### 4. Token Classification (NER/POS)
+Named Entity Recognition and Part-of-Speech tagging.
+*   **Task Type:** `token-classification`
+*   **Features:** Handles label alignment for sub-word tokens automatically.
+
+
+## Data Preparation
+
+The `datasets.py` module handles loading (JSONL, Parquet, CSV, HF Hub) and column standardization. It is built on top of HuggingFace's datasets.
+
+### 1. Column Mapping
+The trainer looks for specific column names. 
+
+| Task | Required Columns | Description |
+| :--- | :--- | :--- |
+| **Classification** | `text`, `label` | Input text and target class/score. |
+| **Pairs (Sim.)** | `text`, `text_pair` | Anchor and Positive/Candidate. |
+| **Triplets** | `text`, `text_pair`, `negative` | Anchor, Positive, Hard Negative. |
+| **MLM** | `text` | Raw text for masking. |
+| **NER** | `tokens`, `labels` | Pre-tokenized words and aligned tags. |
+
+*Note: For Sentence Similarity, if a `label` column is present with floats, the trainer switches to Regression/MSE loss (e.g., for scored Bi-Encoders).*
+
+You can map your custom dataset fields via `DatasetArgs`.
+```
+# Load datasets
+dataset_args = DatasetArgs(
+    data=dataset, # dataset path
+    task_type=task_type, 
+    text_field="question", # maps column 'question' to 'text'
+    text_pair_field="response", # maps column 'response' to 'text_pair'
+    negative_field="semantically_different_response", # maps column 'semantically_different_response' to 'negative'
+    label_field="classification", # maps column 'classification to 'label'
+    test=True # creates a test split, if not already present in the dataset, out of the training set (validation set not affected).
+)
+```  
+Anchor is automatically mapped to `text` and Positive is automatically mapped to `text_pair`. See _standardize_column_names() in `datasets.py` for more information on column mapping. 
+
+### 2. Text Pairs and Chat Template
+
+For certain tasks like text-classification, you may want to classify how two token sequences (text and text_pair) relate to each other.  
+
+For bi-encoders, it is highly recommended to let the tokenizer combine the text and the text_pair rather than aggregating them manually. This ensures that the correct separation token is used.
+
+```
+batch = self.tokenizer(
+    texts,
+    text_pairs,
+    padding="longest",
+    truncation=True,
+    max_length=self.max_length,
+    return_tensors="mlx"
+)
+```
+
+In some cases, you may want to use the chat template that was used to train the model you intend to finetune. For example, LFM2-350M recommends using a chat template.  
+
+If `use_chat_template` is set to True when initializing the training (default False) and if a chat template is available in the tokenizer (do check!), the text and the text_pair values will be combined and text_pair will be set to None.  
+
+You can also force a specific string as separator.  
+
+This is how it works under the hood:
+
+```
+if text_pairs is not None:
+    if getattr(self.tokenizer, "chat_template", None) and self.use_chat_template:
+        # This ensures the model sees exactly what it expects for Q&A
+        formatted_texts = []
+        for prompt, response in zip(texts, text_pairs):
+            messages = [
+                {"role": "user", "content": prompt},
+                {"role": "assistant", "content": response}
+            ]
+            formatted_texts.append(
+                self.tokenizer.apply_chat_template(messages, tokenize=False)
+            )
+        texts = formatted_texts
+        text_pairs = None # Handled by template
+
+    elif self.force_separator is not None:
+        # Use the forced separator for decoder models
+        texts = [
+            f"{t}{self.force_separator}{p}" 
+            for t, p in zip(texts, text_pairs)
+        ]
+        text_pairs = None
+```
+
+See DataCollatorForSequenceClassification in `collators.py` for more information on text_pair handling for text-classification.
+
+
+## Quick Start (Programmatic)
+
+Below is a simplified example of how to set up a training run programmatically.
+
+```python
+from mlx_raclate.utils.utils import load
+from mlx_raclate.tuner.datasets import load_dataset, DatasetArgs
+from mlx_raclate.tuner.trainer import Trainer, TrainingArgs
+
+# 1. Configuration variables
+model_path = "Qwen/Qwen3-Embedding-0.6B"
+dataset_path = "data/wines"
+task_type = "text-classification"
+
+# 2. Load and Prepare Dataset
+dataset_args = DatasetArgs(
+    data=dataset_path, 
+    task_type=task_type,
+    # Optional: override field names if your data isn't standard
+    # text_field="question",
+    # text_pair_field="response",
+    # label_field="classification"
+)
+train_ds, valid_ds, test_ds, id2label, label2id = load_dataset(dataset_args)
+
+# 3. Load Model and Tokenizer
+# Pass label mappings to model config for classification tasks
+model_config = {"id2label": id2label, "label2id": label2id} if id2label else {}
+
+model, tokenizer = load(
+    model_path, 
+    model_config=model_config, 
+    pipeline=task_type,
+    train=True
+)
+
+# 4. Define Training Arguments
+args = TrainingArgs(
+    output_dir="outputs/my_run",
+    batch_size=4,
+    gradient_accumulation_steps=4,
+    learning_rate=1e-5,
+    num_train_epochs=3,
+    lr_scheduler_type="cosine_decay",
+    warmup_ratio=0.03,
+    save_steps=500,
+    logging_steps=10,
+    max_length=2048,
+    freeze_embeddings=False
+)
+
+# 5. Initialize Trainer
+trainer = Trainer(
+    model=model,
+    tokenizer=tokenizer,
+    task_type=task_type,
+    training_args=args,
+    train_dataset=train_ds,
+    eval_dataset=valid_ds,
+    label2id=label2id,
+    # For decoder models doing classification on pairs:
+    use_chat_template=False 
+)
+
+# 6. Run Training
+trainer.train()
+
+# 7. Evaluate on Test Set (Optional)
+if test_ds:
+    trainer.test(test_ds)
+```
+
+## CLI usage 
+
+An example of CLI tool including **all** parameters to train a model is available in `mlx_raclate.utils.train.py`.  
+
+WARNING : this example includes default values that override the default values of the DatasetArgs, TrainingArgs and Trainer classes presented below.
+
+## API Reference
+
+### DatasetArgs
+
+Used to configure how data is loaded and mapped.
+
+| Parameter | Type | Default | Description |
+| :--- | :--- | :--- | :--- |
+| `data` | `str` | *Required* | Local path or HF identifier of the dataset. |
+| `task_type` | `str` | *Required* | The type of task (e.g., `text-classification`). |
+| `text_field` | `str` | `None` | Name of the text input column. |
+| `text_pair_field`| `str` | `None` | Name of the second text input column (for pairs). |
+| `label_field` | `str` | `None` | Name of the label/target column. |
+| `negative_field`| `str` | `None` | Name of the negative samples column. |
+| `test` | `bool` | `False` | If True, creates a test split from the training set if one doesn't exist. |   
+
+Note : use load_dataset("dataset_path") from `datasets.py` to fetch the dataset splits and the label2id dictionary.
+
+### TrainingArgs
+
+Controls the hyperparameters and runtime configuration.
+
+#### Hyperparameters
+| Parameter | Default | Description |
+| :--- | :--- | :--- |
+| `batch_size` | `2` | The physical batch size per device/step. Reduce to 1 on Macbooks with limited RAM if training on long context. |
+| `gradient_accumulation_steps` | `8` | Number of steps to accumulate gradients before updating weights. |
+| `num_train_epochs` | `2` | Total number of training epochs. |
+| `max_length` | `512` | Max sequence length. If `None`, uses model's default config. |
+| `freeze_embeddings` | `False` | If `True`, freezes the embedding layer to save memory/compute. |
+
+#### Optimizer & Scheduler
+| Parameter | Default | Description |
+| :--- | :--- | :--- |
+| `learning_rate` | `3e-5` | Initial learning rate (Peak LR). |
+| `weight_decay` | `0.01` | Weight decay factor for AdamW. |
+| `lr_scheduler_type` | `"constant"` | Scheduler type: `"cosine_decay"`, `"linear_schedule"`, or `"constant"`. |
+| `min_lr` | `0.0` | Minimum learning rate at the end of the schedule. |
+| `warmup_ratio` | `0.0` | Ratio of total training steps used for warmup. |
+| `warmup_steps` | `0` | Absolute number of warmup steps (overrides `warmup_ratio` if set). |
+| `max_grad_norm` | `1.0` | Gradient clipping threshold. |
+
+#### Checkpointing & Logging
+| Parameter | Default | Description |
+| :--- | :--- | :--- |
+| `output_dir` | `None` | Directory to save checkpoints and logs. Defaults to a timestamped folder. |
+| `save_steps` | `1000` | Frequency of saving model checkpoints (in steps). |
+| `logging_steps` | `16` | Frequency of logging metrics to console/files. |
+| `eval_batch_size` | `4` | Batch size used during evaluation/testing. |
+| `resume_from_step`| `0` | Step to resume training from. If this is after the last warmup step (either declared or calculated via warmup_ratio), warmup will be ignored. |  
+
+Gradient checkpointing is enabled by default due to RAM constraints of consumer hardware.
+
+### Model Config
+
+When loading a pretrained model, you can create a model_config dictionary with new parameters and pass it to the load() function. Common examples :
+
+| Parameter | Type | Default | Description |
+| :--- | :--- | :--- | :--- |
+| `is_regression` | `bool` | `False` | For text-classification tasks, whether the classification is a regression |
+| `use_late_interaction`| `bool` | `False` | For sentence similarity tasks, whether late interaction (MaxSim) should be used instead of cosine similarity |
+
+## Trainer
+
+The main class that orchestrates the training.
+
+**Constructor Parameters:**
+
+*   **`model`**: The loaded MLX model.
+*   **`tokenizer`**: The loaded tokenizer. If you want to use a chat template, make sure that the tokenizer includes the chat template. If not, add it manually before instantiating the Trainer.
+*   **`task_type`**: String identifier for the pipeline (e.g., "text-classification").
+*   **`training_args`**: Instance of `TrainingArgs`.
+*   **`train_dataset`**: The processed training dataset.
+*   **`eval_dataset`**: (Optional) The processed validation dataset.
+*   **`label2id`**: (Optional) Dictionary mapping labels to IDs (required for classification metrics).
+*   **`use_chat_template`** *(bool)*: If `True`, applies the tokenizer's chat template to inputs. Useful for decoder models (like Qwen/Llama) performing classification on text pairs.
+*   **`force_separator`** (Optional *str*): If not using a chat template, this string is used to join text pairs for decoder models.
+*   **`optimizer`** (Optional *mlx.optimizer*): If no optimizer is passed, AdamW will be used with the hyper parameters set in TrainingArgs 
+
+**Methods:**
+
+*   `train()`: Starts the training loop.
+*   `test(dataset)`: Runs evaluation on the provided dataset.
+

mlx_raclate/tuner/collators.py

@@ -0,0 +1,291 @@
+import mlx.core as mx
+import numpy as np
+from typing import List, Dict, Any, Optional, Union
+from dataclasses import dataclass
+
+@dataclass
+class DataCollator:
+    tokenizer: Any
+    max_length: int = 512
+    
+    def __call__(self, features: Dict[str, List[Any]]) -> Dict[str, mx.array]:
+        raise NotImplementedError
+
+@dataclass
+class DataCollatorForSequenceClassification(DataCollator):
+    """
+    Handles tokenization and padding for classification tasks.
+    """
+    use_chat_template: bool = False# Whether to use chat templates for decoder models
+    force_separator: Optional[str] = None # If set, forces this separator between text pairs
+    default_decoder_separator: str = "\n" # Used for decoder models when concatenating text pairs
+    label2id: Optional[Dict[str, int]] = None
+
+    def __call__(self, features: Dict[str, List[Any]]) -> Dict[str, mx.array]:
+        texts = features.get("text")
+        text_pairs = features.get("text_pair", None)
+
+        if text_pairs is not None:
+            if getattr(self.tokenizer, "chat_template", None) and self.use_chat_template:
+                # This ensures the model sees exactly what it expects for Q&A
+                formatted_texts = []
+                for prompt, response in zip(texts, text_pairs):
+                    messages = [
+                        {"role": "user", "content": prompt},
+                        {"role": "assistant", "content": response}
+                    ]
+                    formatted_texts.append(
+                        self.tokenizer.apply_chat_template(messages, tokenize=False)
+                    )
+                texts = formatted_texts
+                text_pairs = None # Handled by template
+
+            elif self.force_separator is not None:
+                # Use the forced separator for decoder models
+                texts = [
+                    f"{t}{self.force_separator}{p}" 
+                    for t, p in zip(texts, text_pairs)
+                ]
+                text_pairs = None
+
+            else :
+                # Check if tokenizer has a standard separator (Like [SEP] in BERT)
+                # Qwen tokenizer often has sep_token as None or same as EOS
+                has_sep_token = getattr(self.tokenizer, "sep_token", None) is not None
+                
+                if not has_sep_token or self.tokenizer.sep_token == self.tokenizer.eos_token:
+                    texts = [
+                        f"{t}{self.default_decoder_separator}{p}" 
+                        for t, p in zip(texts, text_pairs)
+                    ]
+                    # Set pairs to None so tokenizer treats it as a single string
+                    text_pairs = None 
+
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        batch = self.tokenizer(
+            texts,
+            text_pairs,
+            padding="longest",
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="mlx"
+        )
+        
+        if "label" in features:
+            labels = features["label"]
+            # On-the-fly String to ID conversion
+            if self.label2id and len(labels) > 0 and isinstance(labels[0], str):
+                labels = [self.label2id.get(l, -1) for l in labels] # Default to -1 if missing
+            
+            # Detect regression (float) vs classification (int)
+            if len(labels) > 0 and isinstance(labels[0], float):
+                dtype = mx.float32
+            else:
+                dtype = mx.int32
+
+            batch["labels"] = mx.array(labels, dtype=dtype)
+            
+        return dict(batch)
+    
+@dataclass
+class DataCollatorForTokenClassification(DataCollator):
+    """
+    Handles tokenization and aligns labels for token classification.
+    """
+    label_pad_token_id: int = -100
+    # Strategy: 'first' (label only first subword), 'all' (label all subwords with same tag)
+    label_all_tokens: bool = False 
+    label2id: Optional[Dict[str, int]] = None
+
+    def __call__(self, features: Dict[str, List[Any]]) -> Dict[str, mx.array]:
+        texts = features["text"] 
+        labels = features["labels"] # Note: usually plural 'labels' list of list
+
+        # SANITY CHECK: The library expects pre-tokenized inputs (List[str])
+        if isinstance(texts[0], str):
+             raise ValueError(
+                 "DataCollatorForTokenClassification expects 'text' to be a list of strings "
+                 "(tokens), not a single string. Please pre-tokenize your dataset."
+             )
+                
+        batch = self.tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="mlx",
+            is_split_into_words=True
+        )
+
+        batch_size, seq_len = batch["input_ids"].shape
+        
+        # Create a numpy buffer filled with the ignore index
+        padded_labels = np.full((batch_size, seq_len), self.label_pad_token_id, dtype=np.int32)
+
+        for i, label_seq in enumerate(labels):
+            # On-the-fly conversion for list of strings (to avoid memory issues with dataset.map)
+            current_labels = label_seq
+            if self.label2id and len(label_seq) > 0 and isinstance(label_seq[0], str):
+                current_labels = [self.label2id.get(l, self.label_pad_token_id) for l in label_seq]
+
+            # word_ids returns a list mapping each token to its original word index
+            # e.g., [None, 0, 1, 1, 2, None] for "[CLS] My name is John [SEP]"
+            word_ids = batch.word_ids(batch_index=i)
+            previous_word_idx = None
+
+            for k, word_idx in enumerate(word_ids):
+                # Skip Special Tokens (None)
+                if word_idx is None:
+                    continue
+                
+                # Safety check: tokenizer truncation might leave word_ids that point to label indices larger than the label list provided.
+                if word_idx >= len(current_labels):
+                    break 
+                
+                if word_idx != previous_word_idx:
+                    padded_labels[i, k] = current_labels[word_idx]
+                else:
+                    # This is a subsequent subword of the same word
+                    if self.label_all_tokens:
+                        padded_labels[i, k] = current_labels[word_idx]
+                    else:
+                        # Standard BERT NER behavior: ignore subsequent subwords
+                        padded_labels[i, k] = self.label_pad_token_id
+                
+                previous_word_idx = word_idx
+
+        batch["labels"] = mx.array(padded_labels, dtype=mx.int32)
+
+        return dict(batch)
+
+@dataclass
+class DataCollatorForMaskedLanguageModeling(DataCollator):
+    """
+    Handles dynamic masking for MLM.
+    """
+    mlm_probability: float = 0.15
+    mask_token_id: Optional[int] = None
+
+    def __call__(self, features: Dict[str, List[Any]]) -> Dict[str, mx.array]:
+        texts = features["text"]
+        
+        batch = self.tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="mlx"
+        )
+        
+        input_ids = batch["input_ids"]
+        
+        # Create Mask
+        probability_matrix = mx.random.uniform(shape=input_ids.shape) < self.mlm_probability
+        
+        # Protect special tokens
+        special_tokens_mask = mx.array([
+            [1 if token_id in [self.tokenizer.cls_token_id, self.tokenizer.sep_token_id, self.tokenizer.pad_token_id]
+             else 0 for token_id in seq]
+            for seq in input_ids.tolist() 
+        ])
+        
+        probability_matrix = mx.where(special_tokens_mask, 0, probability_matrix)
+        
+        # Create labels (-100 for unmasked)
+        labels = mx.where(probability_matrix, input_ids, -100)
+        
+        # Apply masking (80% mask, 10% random, 10% original)
+        random_matrix = mx.random.uniform(shape=input_ids.shape)
+        mask_indices = (probability_matrix) & (random_matrix < 0.8)
+        random_indices = (probability_matrix) & (random_matrix >= 0.8) & (random_matrix < 0.9)
+        
+        # Create masked input
+        masked_inputs = input_ids
+        
+        mask_token_id = self.tokenizer.mask_token_id
+        if mask_token_id is None:
+            if self.mask_token_id is not None:
+                mask_token_id = self.mask_token_id
+            else:
+                 raise ValueError(
+                     "Tokenizer does not have a mask token defined and no mask_token_id provided."
+                 )
+
+        masked_inputs = mx.where(mask_indices, mask_token_id, masked_inputs)
+        random_tokens = mx.random.randint(
+            0, self.tokenizer.vocab_size, 
+            shape=input_ids.shape
+        )
+        
+        # Apply the [MASK] token
+        inputs = mx.where(random_indices, random_tokens, masked_inputs)
+        
+        batch["input_ids"] = inputs
+        batch["labels"] = labels
+        
+        return dict(batch)
+    
+@dataclass
+class DataCollatorForSentenceSimilarity(DataCollator):
+    """
+    Handles data for Bi-Encoder models (Sentence Similarity / Retrieval).
+    Unlike SequenceClassification, this keeps sentences SEPARATE to produce
+    independent embeddings.
+    
+    Expected keys in features (from datasets.py standardization):
+    - 'text': The Anchor / Sentence A
+    - 'text_pair': The Positive / Reference / Sentence B
+    - 'negative' (optional): The Hard Negative / Sentence C
+    - 'label' (optional): Similarity score for Regression
+    """
+    def __call__(self, features: Dict[str, List[Any]]) -> Dict[str, mx.array]:
+        batch = {}
+
+        # Tokenize Anchor (Sentence A) -> 'input_ids'
+        if "text" in features:
+            out_a = self.tokenizer(
+                features["text"],
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="mlx"
+            )
+            batch["input_ids"] = out_a["input_ids"]
+            batch["attention_mask"] = out_a["attention_mask"]
+
+        # Tokenize Reference (Sentence B) -> 'reference_input_ids'
+        if "text_pair" in features:
+            out_b = self.tokenizer(
+                features["text_pair"],
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="mlx"
+            )
+            batch["reference_input_ids"] = out_b["input_ids"]
+            batch["reference_attention_mask"] = out_b["attention_mask"]
+
+        # Tokenize Negative (Sentence C) -> 'negative_input_ids'
+        neg_key = None
+        if "negative" in features: neg_key = "negative"
+        elif "text_negative" in features: neg_key = "text_negative"
+            
+        if neg_key:
+            out_n = self.tokenizer(
+                features[neg_key],
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="mlx"
+            )
+            batch["negative_input_ids"] = out_n["input_ids"]
+            batch["negative_attention_mask"] = out_n["attention_mask"]
+
+        # Handle Scores (for Regression)
+        if "label" in features:
+            # Ensure float32 for regression targets
+            batch["similarity_scores"] = mx.array(features["label"], dtype=mx.float32)
+
+        return batch