rapidfit 0.1.0

Build multi-task classifiers and augment classification datasets with ease

RapidFit

Build multi-task classifiers and augment classification datasets with ease.

Features

• Multi-task Classification: Train classifiers that handle multiple classification tasks simultaneously
• Data Augmentation: Expand and enhance your classification datasets using LLM-based generation
• Flexible Saving: Save data in JSON, JSONL, or CSV formats with incremental or batch saving

Installation

pip install rapidfit

Development Installation

pip install -e .

Quick Start

Data Augmentation

from rapidfit import LLMAugmenter, SaveFormat

# Prepare seed data
seed_data = {
    "sentiment-analysis": [
        {"text": "I love this product!", "label": "positive"},
        {"text": "Terrible experience.", "label": "negative"},
        {"text": "It's okay, nothing special.", "label": "neutral"},
    ],
    "emotion-analysis": [
        {"text": "This makes me so happy!", "label": "joy"},
        {"text": "I can't believe they did this.", "label": "anger"},
        {"text": "I miss the old days.", "label": "sadness"},
    ],
}

# Initialize augmenter
augmenter = LLMAugmenter(
    api_key="your-openai-api-key",
    base_url=None,  # Optional: custom API endpoint
    model_id="gpt-4.1-mini",  # Optional: model to use
    max_samples_per_task=128,  # Optional: max samples per task
    batch_size=8,  # Optional: samples per generation batch
    max_temperature=0.9,  # Optional: max temperature for sampling
    save_path="./saved",  # Optional: output directory
    save_format=SaveFormat.JSON,  # Optional: json, jsonl, or csv
    save_separated=False,  # Optional: separate file per task
    save_incremental=True,  # Optional: save while generating
)

# Augment dataset
augmented_data = augmenter.augment(seed_data)

Save Options

Parameter	Default	Description
save_path	"./saved"	Directory for output files
save_format	json	Output format: json, jsonl, csv
save_separated	False	Create separate file for each task
save_incremental	True	Save progressively during generation

Custom Augmenter

Extend BaseAugmenter to create custom augmentation strategies:

from rapidfit import BaseAugmenter, SeedData

class MyAugmenter(BaseAugmenter):
    def augment(self, seed_data: SeedData) -> SeedData:
        # Your augmentation logic here
        return seed_data

Running

Install Dependencies

pip install -e .

Run Augmentation

export OPENAI_API_KEY="your-api-key"
python examples/test_augmentation.py

Optional environment variables:

• OPENAI_BASE_URL - Custom API endpoint
• OPENAI_MODEL_ID - Model to use (default: gpt-4.1-mini)

Output saves to ./saved/ directory.

Classification

Training a Multihead Classifier

from rapidfit import MultiheadClassifier

# Prepare data (or load from augmented files)
seed_data = {
    "sentiment": [
        {"text": "I love this!", "label": "positive"},
        {"text": "Terrible.", "label": "negative"},
    ],
    "emotion": [
        {"text": "So happy!", "label": "joy"},
        {"text": "I'm angry.", "label": "anger"},
    ],
}

# Create classifier with custom config
classifier = MultiheadClassifier({
    "model_name": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    "batch_size": 16,
    "epochs": 10,
    "freeze_epochs": 3,
    "patience": 3,
})

# Train
classifier.train(seed_data)

# Save model
classifier.save("./model")

# Predict
predictions = classifier.predict(["Great product!"], task="sentiment")
print(predictions)  # [{"label": "positive", "confidence": 0.95}]

# Predict all tasks at once
all_preds = classifier.predict_all_tasks(["Great product!"])

Load a Trained Model

from rapidfit import MultiheadClassifier

classifier = MultiheadClassifier()
classifier.load("./model")

predictions = classifier.predict(["Test text"], task="sentiment")

Training Configuration

Parameter	Default	Description
model_name	paraphrase-multilingual-MiniLM-L12-v2	HuggingFace model
batch_size	16	Training batch size
epochs	10	Fine-tuning epochs
freeze_epochs	3	Epochs with frozen encoder
learning_rate	2e-5	Learning rate for fine-tuning
patience	3	Early stopping patience
dropout_rate	0.2	Dropout rate
label_smoothing	0.1	Label smoothing factor
use_class_weights	True	Handle class imbalance
test_size	0.1	Test split ratio
val_size	0.1	Validation split ratio

Run Training Example

# First, generate augmented data
python examples/test_augmentation.py

# Then train classifier
python examples/train_classifier.py

Custom Classifier

Extend BaseClassifier to create custom classification strategies:

from rapidfit import BaseClassifier
from rapidfit.types import AugmentResult, Prediction, SeedData

class MyClassifier(BaseClassifier):
    def train(self, data: SeedData | AugmentResult) -> None:
        samples = self._resolve_data(data)
        # Training logic here

    def predict(self, texts: list[str], task: str) -> list[Prediction]:
        # Prediction logic here
        return [{"label": "class", "confidence": 0.95}]

    def save(self, path):
        # Save model

    def load(self, path):
        # Load model

Available Classifier Types

Type	Description
MULTIHEAD	Shared encoder with task-specific classification heads

License

MIT