Swarmauri Logo

Swarmauri Embedding MLM

Trainable embedding provider that fine-tunes a Hugging Face masked language model (MLM) end-to-end so Swarmauri agents can produce contextual document vectors without leaving the framework.

Features

Wraps any Hugging Face masked language model (embedding_name) behind the Swarmauri EmbeddingBase interface.
Supports optional vocabulary expansion via add_new_tokens before fine-tuning to capture domain-specific terminology.
Handles end-to-end fine-tuning with masking, AdamW optimization, and GPU/CPU selection based on availability.
Exposes pooling utilities (transform, infer_vector) that average the last hidden state to yield dense vectors ready for downstream retrieval or clustering.
Provides save_model/load_model helpers so trained weights and tokenizers can be persisted and reloaded across workers.

Prerequisites

Python 3.10 or newer.
PyTorch with CUDA support if you plan to train on GPU (the class falls back to CPU automatically).
Access to the Hugging Face model hub for downloading embedding_name. Set HF_HOME, proxies, or tokens if your environment requires authentication.
Enough disk space to cache the chosen MLM (e.g., bert-base-uncased ~420 MB).

Installation

# pip
pip install swarmauri_embedding_mlm

# poetry
poetry add swarmauri_embedding_mlm

# uv (pyproject-based projects)
uv add swarmauri_embedding_mlm

Quickstart: Fine-tune and Embed Documents

from swarmauri_embedding_mlm import MlmEmbedding

docs = [
    "Swarmauri SDK ships modular agents.",
    "Masked language models produce contextual embeddings.",
]

embedder = MlmEmbedding(
    embedding_name="distilbert-base-uncased",
    batch_size=16,
    learning_rate=3e-5,
)

# One epoch of MLM fine-tuning on your corpus
embedder.fit(docs)

# Generate vectors for downstream tasks
vectors = embedder.transform([
    "Agents coordinate through shared memory",
    "Fine-tuning improves domain recall",
])

for v in vectors:
    print(len(v.value), v.value[:4])  # dimension and preview

# Single-text inference helper
query_vector = embedder.infer_vector("How do masked models compute embeddings?")

Expanding the Vocabulary

Set add_new_tokens=True to capture domain-specific terms before training. New tokens are identified via simple whitespace tokenization and appended to the tokenizer before the first epoch.

from swarmauri_embedding_mlm import MlmEmbedding

domain_docs = [
    "Neo4j graph embeddings power fraud detection",
    "Qdrant supports hybrid sparse-dense search",
]

embedder = MlmEmbedding(add_new_tokens=True)
embedder.fit(domain_docs)

# Inspect the tokenizer to confirm additions
print(f"Vocabulary size: {len(embedder.extract_features())}")

Persisting and Reloading Models

from pathlib import Path
from swarmauri_embedding_mlm import MlmEmbedding

save_dir = Path("models/mlm-distilbert")

embedder = MlmEmbedding()
embedder.fit(["short corpus", "to warm up the model"])
embedder.save_model(save_dir.as_posix())

# Later or on another machine
restored = MlmEmbedding()
restored.load_model(save_dir.as_posix())

embedding = restored.infer_vector("Reuse the trained weights instantly")

Operational Tips

Batch and sequence length drive GPU memory usage; reduce batch_size or max_length in tokenizer calls when running on constrained hardware.
fit_transform runs a full fine-tuning pass and immediately returns embeddings—useful for one-off adaptation jobs.
When training on large corpora, stream documents from a generator, chunk them, or wrap the .fit call in your own epoch loop.
Run extract_features() to audit the tokenizer vocabulary (helpful when debugging domain token coverage).
Combine the generated vectors with Swarmauri vector stores (Redis, Qdrant, etc.) to build end-to-end retrieval pipelines.

Want to help?

If you want to contribute to swarmauri-sdk, read up on our guidelines for contributing that will help you get started.

Metadata

Release files for swarmauri_embedding_mlm 0.8.2

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for swarmauri_embedding_mlm 0.8.2
File	Size	Uploaded
swarmauri_embedding_mlm-0.8.2.tar.gz	9.9 kB	Mar 24, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for swarmauri_embedding_mlm 0.8.2
File	Interpreter	ABI	Platform	Reset
swarmauri_embedding_mlm-0.8.2-py3-none-any.whl 11.0 kB Mar 24, 2026	Python 3	none	any	Details

Total release size: 20.9 kB

Release files / swarmauri_embedding_mlm-0.8.2.tar.gz

Download URL	swarmauri_embedding_mlm-0.8.2.tar.gz
Size	9.9 kB
Tags	Source
SHA-256 checksum How to use checksums	`a2ae6f5523754ff6520527acace8a0297d01aafa9c92edc4f105e930090495eb`
BLAKE2b-256 checksum How to use checksums	`0d9311641a246be463cd5039c4e998dfb419c7cb50d1509ea4b857d5eaedb384`
Upload date	Mar 24, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.11.0 {"installer":{"name":"uv","version":"0.11.0","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}`

Release files / swarmauri_embedding_mlm-0.8.2-py3-none-any.whl

Download URL	swarmauri_embedding_mlm-0.8.2-py3-none-any.whl
Size	11.0 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`b238e904ba94a4c1c3a17e90b149b47def90c09c2fd5526b49c514d20a290cd3`
BLAKE2b-256 checksum How to use checksums	`5db7efda631a9f855903ba121421c57a6ed3a7f3db78c7c01330d1917b4655e3`
Upload date	Mar 24, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.11.0 {"installer":{"name":"uv","version":"0.11.0","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}`