mlx-raclate 0.1.0b1__py3-none-any.whl

mlx_raclate-0.1.0b1.dist-info/METADATA

@@ -0,0 +1,216 @@
+Metadata-Version: 2.4
+Name: mlx-raclate
+Version: 0.1.0b1
+Summary: Raclate is a python library to run and train models for Retrieval and Classification, built on top of MLX.
+Author-email: pappitti <pap@pitti.io>
+License-File: LICENSE
+Requires-Python: >=3.13
+Requires-Dist: datasets>=4.4.1
+Requires-Dist: fastapi>=0.121.3
+Requires-Dist: huggingface-hub>=0.36.0
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: mlx>=0.29.0
+Requires-Dist: transformers>=4.57.1
+Requires-Dist: uvicorn>=0.38.0
+Description-Content-Type: text/markdown
+
+# RACLATE (MLX)
+
+**R**etrieval **A**nd **C**lassification including **LATE** interaction on Apple Silicon.
+
+`mlx-raclate` is a versatile library built on Apple's [MLX](https://github.com/ml-explore/mlx) framework. It provides a unified interface to **train** and **run** classifiers and embedding models - including ModernBERT and Late Interaction (ColBERT-style) models - natively on macOS.
+
+> **Note:** This project evolved from `modernbert-mlx` to support a wider range of architectures and tasks. It is currently feature-complete but in an early release stage; bugs may occur.
+
+## Key Features
+
+*   **Apple Silicon Native:** Fully optimized for M-series chips using MLX.
+*   **Unified Pipeline:** A single interface to load and run Masked LM, Text Classification, and Sentence Similarity tasks.
+*   **Late Interaction Support:** First-class support for **MaxSim** (ColBERT-style) retrieval, particularly with LFM2 and ModernBERT architectures.
+*   **Full Fine-Tuning:** specialized trainer for fine-tuning small-to-mid-sized models (ModernBERT, Qwen2.5/3, LFM2, Gemma) on local hardware.
+
+## Installation
+
+Install via `uv` or `pip`:
+
+```bash
+uv add --prerelease=allow mlx-raclate
+# or
+pip install --pre mlx-raclate 
+```
+
+From source:
+
+```bash
+git clone https://github.com/pappitti/mlx-raclate.git
+cd mlx-raclate
+uv sync
+```
+
+## Supported Architectures
+
+`mlx-raclate` supports architectures specifically useful for efficient local retrieval and classification:
+
+*   **ModernBERT**: (e.g., `answerdotai/ModernBERT-base`)
+*   **LFM2**: Liquid Foundation Models (e.g., `LiquidAI/LFM2-350M`, `LiquidAI/LFM2-ColBERT-350M`)
+*   **Qwen3 Embedding**: (e.g., `Qwen/Qwen3-Embedding-0.6B`)
+*   **Gemma3 Embedding**: (e.g., `google/embeddinggemma-300m`)
+*   **T5Gemma Encoder**: stripping out the decoder part of T5Gemma models (e.g, `google/t5gemma-2b-2b-ul2`)
+
+## Inference: Quick Start
+
+The library uses a `pipeline` concept similar to Hugging Face Transformers. You can specify a pipeline manually, or let the loader infer it from the model configuration.  
+If no pipeline is found, the Model class is loaded, which returns normalized embeddings.
+
+### 1. Text Classification
+Supports multi-class, multi-label, and regression tasks.
+
+```python
+from mlx_raclate.utils.utils import load
+import mlx.core as mx
+
+# Load model (pipeline inferred automatically if architecture matches)
+model, tokenizer = load("NousResearch/Minos-v1", pipeline="text-classification")
+
+texts = ["How do I build a bomb?", "What is the capital of France?"]
+
+# Batch tokenize
+inputs = tokenizer._tokenizer(texts, return_tensors="mlx", padding=True, truncation=True)
+
+# Run Inference
+outputs = model(
+    input_ids=inputs['input_ids'], 
+    attention_mask=inputs['attention_mask']
+)
+
+# Get probabilities
+probs = outputs["probabilities"]
+# ... process argmax/topk
+```
+
+### 2. Sentence Similarity (Dense Retrieval)
+2.1 Standard Bi-Encoder approach using Cosine Similarity.
+
+```python
+from mlx_raclate.utils.utils import load
+
+model, tokenizer = load("nomic-ai/modernbert-embed-base", pipeline="sentence-similarity")
+
+queries = ["What is MLX?"]
+docs = ["MLX is an array framework for Apple Silicon."]
+
+# Encode
+q_input = tokenizer._tokenizer(queries, return_tensors="mlx", padding=True)
+d_input = tokenizer._tokenizer(docs, return_tensors="mlx", padding=True)
+
+# Forward pass calculates similarity matrix automatically
+outputs = model(
+    input_ids=q_input['input_ids'],
+    reference_input_ids=d_input['input_ids'],
+    attention_mask=q_input['attention_mask'],
+    reference_attention_mask=d_input['attention_mask']
+)
+
+print(outputs['similarities']) # Cosine similarity matrix
+```
+
+2.2. Late Interaction (ColBERT / MaxSim)
+By enabling `use_late_interaction`, the model computes **MaxSim** scores (interaction between all token embeddings) instead of standard Cosine similarity of pooled embeddings.
+
+This is ideal for models like **LFM2-ColBERT**, but it works with any model.
+
+```python
+from mlx_raclate.utils.utils import load
+
+# Load a ColBERT-style model
+model, tokenizer = load(
+    "LiquidAI/LFM2-ColBERT-350M", 
+    pipeline="sentence-similarity",
+    model_config={"use_late_interaction": True} # <--- Enables MaxSim
+)
+
+queries = ["Who creates liquid neural networks?"]
+docs = ["Liquid AI is a company founded by researchers from MIT..."]
+
+# Tokenize
+q_input = tokenizer._tokenizer(queries, return_tensors="mlx", padding=True)
+d_input = tokenizer._tokenizer(docs, return_tensors="mlx", padding=True)
+
+# The model keeps embeddings unpooled and computes MaxSim
+outputs = model(
+    input_ids=q_input['input_ids'],
+    reference_input_ids=d_input['input_ids'],
+    attention_mask=q_input['attention_mask'],
+    reference_attention_mask=d_input['attention_mask']
+)
+
+print("MaxSim Scores:", outputs['similarities'])
+```
+
+## Pipelines Reference
+
+When using `load()`, the `pipeline` argument determines the class and return types. If not provided, `mlx-raclate` attempts to infer it from the `config.json`.
+
+| Pipeline | Class | Output | Use Case |
+| :--- | :--- | :--- | :--- |
+| `embeddings` | `Model` | Raw Embeddings | Feature extraction |
+| `text-classification` | `ModelForSequenceClassification` | Logits/Probs | Sentiment, Intent, Regression |
+| `sentence-similarity` | `ModelForSentenceSimilarity` | Embeddings & Similarity | Semantic Search, RAG |
+| `sentence-transformers` | `ModelForSentenceTransformers` | Embeddings & Similarity | Same as `sentence-similarity` but  different sanitization strategy for Sentence Transformers weights |
+| `masked-lm` | `ModelForMaskedLM` | Token Logits | Domain adaptation, MLM training |
+| `token-classification` | `ModelForTokenClassification` | Token Logits | NER tasks | 
+| `zero-shot-classification` | `ModelForMaskedLM` | Token Logits | Implementation of [this AnswerAI paper](https://arxiv.org/html/2502.03793v2) |
+
+Detailed code for each pipeline is available in the `test` directory of this repository. See `tests/inference_examples`.
+
+## Server
+
+`mlx-raclate` includes a FastAPI server for classifier inference. See `mlx_raclate.utils.server`
+
+## Training (Tuner)
+
+`mlx-raclate` includes a robust training engine specifically designed for fine-tuning these architectures on Apple Silicon.
+
+It supports:
+*   **Full Fine-tuning** (LoRA is not currently supported/needed for these model sizes).
+*   **Tasks:** Text Classification, Sentence Similarity (Bi-Encoder & Late Interaction), and Masked LM.
+*   **Efficiency:** Gradient Accumulation, Gradient Checkpointing, and Smart Collation.
+
+For detailed training documentation, supported datasets, and CLI usage, please see [TUNER.md](src/mlx_raclate/tuner/TUNER.md).
+
+### Quick Training Snippet
+
+```python
+from mlx_raclate.tuner.trainer import Trainer, TrainingArgs
+from mlx_raclate.utils.utils import load
+
+# Load model
+model, tokenizer = load("Qwen/Qwen3-Embedding-0.6B", pipeline="text-classification", train=True)
+
+# Define Args
+args = TrainingArgs(
+    output_dir="outputs/my_classifier",
+    learning_rate=1e-5,
+    num_train_epochs=3,
+    batch_size=4
+)
+
+# Initialize Trainer
+trainer = Trainer(
+    model=model,
+    tokenizer=tokenizer,
+    task_type="text-classification",
+    training_args=args,
+    train_dataset=train_dataset, # See TUNER.md for dataset formatting
+    eval_dataset=eval_dataset
+)
+
+trainer.train()
+```
+
+## Acknowledgements
+
+*   [MLX](https://github.com/ml-explore/mlx) team for the framework.
+*   [Transformers](https://github.com/huggingface/transformers) for the configuration standards.
+*   [MLX-Embeddings](https://github.com/Blaizzy/mlx-embeddings) for inspiration on broader embeddings architecture. MLX-Raclate focuses on longer-context models but you should definitely look there for BERT, XLM_RoBERTa and image embeddings.
+*   [PyLate](https://github.com/lightonai/pylate) for inspiration on Late Interaction mechanics.

mlx_raclate-0.1.0b1.dist-info/RECORD

@@ -0,0 +1,25 @@
+mlx_raclate/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+mlx_raclate/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlx_raclate/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlx_raclate/models/base.py,sha256=0V8CfjWIoeZVwd7oWV5M3lif4k9P2R-DjVYaL5yqaSE,8794
+mlx_raclate/models/gemma3_text.py,sha256=daSVYlKs4OITUuuI5KJPVSSkjfske3YLlYHjQ9DNMHU,32519
+mlx_raclate/models/lfm2.py,sha256=o_TmjnZR3o35YG-3n3u2no4cVUfUvTOfQElAnC4S5yA,22885
+mlx_raclate/models/modernbert.py,sha256=esMexSVnjvqz8-bcdamXSw0-pAP_lL1wws7Gk9OvLqM,33486
+mlx_raclate/models/qwen3.py,sha256=4xiJyXs9vvIF_NBini-eYscesYgADZa_KEzZPerRaDY,20339
+mlx_raclate/models/t5gemma_encoder.py,sha256=Q24jMAYtSB6DYwmF1XJ5lxQsO0FACH836jwTHIF_Co0,30377
+mlx_raclate/tuner/TUNER.md,sha256=w_FgVIrV2U3ob_bwJwmGKDBWzw_7bS14wozAosas3x0,14983
+mlx_raclate/tuner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlx_raclate/tuner/collators.py,sha256=eELZWrwYTDUXaJTj4Go2pwnKhlS6yzdS8RNLvFsOhbU,11462
+mlx_raclate/tuner/datasets.py,sha256=xnpTCTqpX8tIrb71E47E055wzxUSBDS3-uCGk3WAoYE,11212
+mlx_raclate/tuner/model_card_utils.py,sha256=XeYia6OMCdZmt1k7NpLTfnpnlODTk3ztCz-2PjCYRMU,6466
+mlx_raclate/tuner/trainer.py,sha256=WU9RkYQxhCtjhaLKIzOAhDOQQ6XHGgqFnuIJfUvZHPc,26734
+mlx_raclate/tuner/utils.py,sha256=Vo95eX0NWTuRXv4pshRp-VTcnaYLtHprRXiR0snndQY,9879
+mlx_raclate/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlx_raclate/utils/server.py,sha256=2BnwGyG0YnJADqMrsvpN5DPyTEANozeLaaOACpbICfU,15541
+mlx_raclate/utils/tokenizer_utils.py,sha256=vrk5tVS0GfeP8OKgGGyI2eS4J2wOcxaBy5jqPRSXs3c,11415
+mlx_raclate/utils/train.py,sha256=p7jTCTQXN7I5yc_LkdzOY84YaYaVTlhfpjGP3zq8k_U,12929
+mlx_raclate/utils/utils.py,sha256=2SszegS5ppmREfIdgr7xKweY1R5u00PoyYnaNI1BG2w,23526
+mlx_raclate-0.1.0b1.dist-info/METADATA,sha256=SvOVI0DjIQ7B4FTDbES61lqUQF_pSBEb2rlO20HODNc,8344
+mlx_raclate-0.1.0b1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mlx_raclate-0.1.0b1.dist-info/licenses/LICENSE,sha256=GEfg4GmBQu1DR8FEGp-oHI-93USx2LvNXjZH-ZF1nX8,1035
+mlx_raclate-0.1.0b1.dist-info/RECORD,,

mlx_raclate-0.1.0b1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

mlx_raclate-0.1.0b1.dist-info/licenses/LICENSE

@@ -0,0 +1,19 @@
+MIT License
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.