bit_transformer/__init__.py

"""BitTransformerLM: Bit-native transformer with reversible layers and telemetry."""

# Core model components
from .model import (
    BitTransformerLM,
    PositionalEncoding,
    ReversibleLoggingTransformerEncoderLayer,
    diffusion_inference,
    example_training_step,
    example_usage,
    infer_long_sequence,
)

# I/O and data processing
from .bit_io import bits_to_text, infer_text, text_to_bits
from .compression import (
    compress_bits,
    decompress_bits,
    model_output_decompress,
    pack_bits,
    unpack_bits,
)
from .parity import enforce_parity

# Training and optimization
from .optimization import adjust_learning_rate, configure_optimizer
from .training import train_loop

# Model scaling and distillation
from .collapse import collapse_submodel, save_distilled_model
from .distil import TelemetryLog, distill_step
from .scale import expand_model

# Distributed computing
from .distributed import make_pipeline, wrap_fsdp

# Quantization support
from .quantization import convert_qat_fx, prepare_qat_fx, quantize_dynamic

# Safety and monitoring
from .safety import demo_hil_safety, hil_safe_inference, safe_sample_with_retry
from .telemetry import TelemetrySynthesizer, detect_metric_drift

# Configuration management
from .config import (
    DataConfig,
    ExperimentConfig,
    ModelConfig,
    SafetyConfig,
    TrainingConfig,
    get_config_from_env,
    get_large_config,
    get_medium_config,
    get_small_config,
)

# Command-line interface
from .cli import dashboard_cli, infer_cli, train_cli
from .cli_standards import BitTransformerCLI

# Visualization and utilities
from .dashboard import plot_telemetry
from .dashboard_app import run_dashboard
from .hf_checkpoint import download_checkpoint, hf_login, save_checkpoint
from .torch_utils import cpu_autocast
from .utils import load_model, save_model, set_dropout

__all__ = [
    # Core model components
    "BitTransformerLM",
    "PositionalEncoding",
    "ReversibleLoggingTransformerEncoderLayer",
    "diffusion_inference",
    "example_training_step",
    "example_usage",
    "infer_long_sequence",

    # I/O and data processing
    "bits_to_text",
    "compress_bits",
    "decompress_bits",
    "enforce_parity",
    "infer_text",
    "model_output_decompress",
    "pack_bits",
    "text_to_bits",
    "unpack_bits",

    # Training and optimization
    "adjust_learning_rate",
    "configure_optimizer",
    "train_loop",

    # Model scaling and distillation
    "collapse_submodel",
    "distill_step",
    "expand_model",
    "save_distilled_model",
    "TelemetryLog",

    # Distributed computing
    "make_pipeline",
    "wrap_fsdp",

    # Quantization support
    "convert_qat_fx",
    "prepare_qat_fx",
    "quantize_dynamic",

    # Safety and monitoring
    "demo_hil_safety",
    "detect_metric_drift",
    "hil_safe_inference",
    "safe_sample_with_retry",
    "TelemetrySynthesizer",

    # Configuration management
    "DataConfig",
    "ExperimentConfig",
    "get_config_from_env",
    "get_large_config",
    "get_medium_config",
    "get_small_config",
    "ModelConfig",
    "SafetyConfig",
    "TrainingConfig",

    # Command-line interface
    "BitTransformerCLI",
    "dashboard_cli",
    "infer_cli",
    "train_cli",

    # Visualization and utilities
    "cpu_autocast",
    "download_checkpoint",
    "hf_login",
    "load_model",
    "plot_telemetry",
    "run_dashboard",
    "save_checkpoint",
    "save_model",
    "set_dropout",
]