composennent 0.4.4

Composable neural network components for building models in PyTorch.

composennent

Composable neural network components for building models in PyTorch.

Composennent provides modular, reusable building blocks for constructing transformer-based models. Train GPT, BERT, and other architectures with minimal code.

Features

• 🧩 Modular Components: Encoder, Decoder, Attention blocks that compose together
• 🚀 Built-in Training: Pre-training and fine-tuning with a single method call
• 📝 Multiple Architectures: GPT, BERT, Seq2Seq support out of the box
• 🔧 Tokenizer Support: WordPiece and SentencePiece tokenizers included
• ⚡ Mixed Precision: Automatic mixed precision (AMP) support
• 🎯 Instruction Tuning: Fine-tune models on instruction datasets (Alpaca format)

Installation

pip install composennent

For tokenizer support:

pip install composennent[tokenizers]

For development:

pip install composennent[dev]

Quick Start

Pre-train a GPT Model

import torch
from composennent.models import GPT
from composennent.nlp.tokenizers import SentencePieceTokenizer

# Create model
model = GPT(
    vocab_size=32000,
    latent_dim=512,
    num_heads=8,
    num_layers=6,
    max_seq_len=512,
)

# Load tokenizer
tokenizer = SentencePieceTokenizer.from_pretrained("tokenizer.model")

# Pre-train
texts = ["Your training data here...", ...]
model.pretrain(
    texts=texts,
    tokenizer=tokenizer,
    epochs=3,
    batch_size=16,
    device="cuda",
)

# Save
model.save("my_model.pt")

Fine-tune on Instructions

# Load pre-trained model
model = GPT.load("my_model.pt", device="cuda")

# Instruction data (Alpaca format)
instruction_data = [
    {
        "instruction": "What is the capital of France?",
        "input": "",
        "output": "The capital of France is Paris."
    },
    # ... more examples
]

# Fine-tune
model.fine_tune(
    data=instruction_data,
    tokenizer=tokenizer,
    epochs=2,
    lr=5e-5,
    mask_prompt=True,  # Only compute loss on outputs
)

Generate Text

prompt = tokenizer.encode("What is")
generated = model.generate(
    input_ids=prompt,
    max_length=100,
    temperature=0.8,
)
print(tokenizer.decode(generated[0].tolist()))

Modules

Module	Description
composennent.modules	Core building blocks (Encoder, Decoder, Block)
composennent.modules.attention	Attention mechanisms and masks
composennent.models	GPT, BERT, and other transformer models
composennent.nlp.tokenizers	WordPiece and SentencePiece tokenizers
composennent.trainer	Training utilities and trainer classes
composennent.modules.experts	Mixture of Experts components
composennent.vision	Vision transformer components
composennent.utils	Utility functions

Training API

For more control over training, use the trainer classes directly:

from composennent.trainer import CausalLMTrainer, train

# Option 1: Use the train() convenience function
train(model, texts, tokenizer, model_type="causal_lm", epochs=5)

# Option 2: Use trainer class directly
trainer = CausalLMTrainer(model, tokenizer, device="cuda")
trainer.train(texts, epochs=5, batch_size=16)
trainer.save_checkpoint("checkpoint.pt")

Available trainers:

• CausalLMTrainer - GPT-style next-token prediction
• MaskedLMTrainer - BERT-style masked language modeling
• Seq2SeqTrainer - Encoder-decoder models
• MultiTaskTrainer - Multi-task learning (MLM + NSP)
• CustomTrainer - Custom loss functions

Requirements

• Python >= 3.8
• PyTorch >= 2.0.0

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

1. Fork the repository
2. Create your feature branch (git checkout -b feature/amazing-feature)
3. Install dev dependencies (pip install -e ".[dev]")
4. Run tests (pytest)
5. Run formatters (black . && ruff check .)
6. Commit your changes (git commit -m 'Add amazing feature')
7. Push to the branch (git push origin feature/amazing-feature)
8. Open a Pull Request

License

MIT License - see LICENSE for details.

Links

• Repository
• PyPI