Pretense

Pretraining sentence transformers with retrieval-oriented objectives.

Pretense supports:

• RetroMAE
• DupMAE
• Condenser
• coCondenser
• Contriever
• supervised pairwise contrastive training
• Multiple Negatives Ranking Loss (MNRL)
• cached MNRL (CMNRL)
• supervised and unsupervised SimCSE

Training uses Hugging Face Transformers and Accelerate, accepts Hugging Face Datasets or ordinary PyTorch datasets, and produces encoders that load in Transformers and Sentence Transformers 5.x or 6.x.

Installation

uv add pretense

Python 3.10 or newer is required. CUDA-enabled PyTorch should be selected using the appropriate uv/PyTorch index for the target system.

Pretense supports PyTorch 2.2 and newer without choosing or replacing a particular CUDA build. Restoring a complete Trainer checkpoint, including optimizer and scheduler state, requires PyTorch 2.6 or newer because Transformers blocks torch.load checkpoint restoration on earlier releases. Weights-only Pretense checkpoints use safetensors and remain loadable on every supported version.

For development from a source checkout, use uv sync --extra dev instead. The core install accepts ordinary Python or PyTorch datasets without installing Hugging Face Datasets. To pass a datasets.Dataset directly or run the bundled examples, install the optional dependency with uv add "pretense[examples]".

Train with the Python SDK

PretenseTrainer follows the Hugging Face Trainer interface:

from transformers import AutoTokenizer

from pretense import (
    MAECollator,
    MethodConfig,
    PretenseTrainer,
    PretenseTrainingArguments,
    load_pretraining_model,
)

model_name = "google-bert/bert-base-uncased"
method = MethodConfig(name="retromae")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = load_pretraining_model(method, model_name)
dataset = [
    {"text": "Dense retrieval maps queries and passages into a shared embedding space."},
    {"text": "Masked autoencoders learn representations by reconstructing corrupted text."},
    {"text": "Sentence transformers produce reusable fixed-size text embeddings."},
    {"text": "In-batch negatives make contrastive learning efficient."},
]

trainer = PretenseTrainer(
    model=model,
    args=PretenseTrainingArguments(
        output_dir="outputs/retromae",
        per_device_train_batch_size=16,
        learning_rate=5e-5,
        num_train_epochs=1,
    ),
    train_dataset=dataset,
    data_collator=MAECollator(tokenizer=tokenizer, text_column="text"),
    processing_class=tokenizer,
)
trainer.train()
trainer.save_model("outputs/retromae/final")

The trainer accepts the same callbacks, optimizers, schedulers, evaluation datasets, and metrics as transformers.Trainer. Ordinary transformers.TrainingArguments are also supported. When the dataset uses the standard columns shown below, the trainer can select the collator automatically; pass a collator explicitly for custom column names or explicit MNRL negatives.

For complete Python workflows, including programmatic models and Sentence Transformers fine-tuning, see the examples.

Training produces regular console and training_log.jsonl metrics and resumable checkpoint-* directories. See training and checkpointing for evaluation, retention, callbacks, experiment trackers, recovery after interruption, and exporting.

Export a trained model explicitly when it is ready for downstream use:

from pretense import export_sentence_transformer
from transformers import AutoModel
from sentence_transformers import SentenceTransformer

export_dir = export_sentence_transformer(
    trainer.model,
    tokenizer,
    "outputs/retromae/sentence-transformers",
)
sentence_model = SentenceTransformer(str(export_dir))
encoder = AutoModel.from_pretrained(export_dir / "0_Transformer")

The Sentence Transformers directory is the canonical export. It includes the complete Hugging Face backbone under 0_Transformer/, so separate copies of the model weights are unnecessary. For a Hub export, load that backbone with AutoModel.from_pretrained(repo_id, subfolder="0_Transformer").

Supported methods

See the method notes for input formats, configuration, pooling behavior, distributed-training caveats, and architecture support. Separate guides cover custom models and FlashAttention.

Development

uv run ruff check .
uv run mypy src/pretense
uv run pytest
uv build --no-sources

Pretense targets objective and architecture parity, not guaranteed reproduction of paper benchmark scores. See the method notes and release documentation.

Attribution

The implementation draws on the following papers and reference projects:

• RetroMAE and DupMAE: Apache-2.0 reference implementation
• Condenser and coCondenser: Apache-2.0 reference implementation
• Contriever: archived reference implementation and Unsupervised Dense Information Retrieval with Contrastive Learning
• Pairwise contrastive loss: Dimensionality Reduction by Learning an Invariant Mapping
• MNRL: Efficient Natural Language Response Suggestion for Smart Reply
• CMNRL: Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup
• SimCSE: SimCSE: Simple Contrastive Learning of Sentence Embeddings and the official implementation

Cite the corresponding paper when publishing results. Pretense is licensed under Apache-2.0.