🚀 Refined BitTransformerLM: Organized codebase with best practices

bit_transformer/__init__.py

@@ -1,23 +1,18 @@
+"""BitTransformerLM: Bit-native transformer with reversible layers and telemetry."""
+
+# Core model components
 from .model import (
-    PositionalEncoding,
     BitTransformerLM,
+    PositionalEncoding,
     ReversibleLoggingTransformerEncoderLayer,
-    example_usage,
+    diffusion_inference,
     example_training_step,
+    example_usage,
     infer_long_sequence,
-    diffusion_inference,
 )
-from .telemetry import TelemetrySynthesizer, detect_metric_drift
-from .dashboard import plot_telemetry
-from .dashboard_app import run_dashboard
-from .collapse import collapse_submodel, save_distilled_model
-from .safety import hil_safe_inference, demo_hil_safety, safe_sample_with_retry
-from .bit_io import (
-    text_to_bits,
-    bits_to_text,
-    infer_text,
-)
-from .parity import enforce_parity
+
+# I/O and data processing
+from .bit_io import bits_to_text, infer_text, text_to_bits
 from .compression import (
     compress_bits,
     decompress_bits,
@@ -25,62 +20,125 @@ from .compression import (
     pack_bits,
     unpack_bits,
 )
-from .distributed import wrap_fsdp, make_pipeline
-from .optimization import configure_optimizer, adjust_learning_rate
+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
-from .distil import distill_step, TelemetryLog
-from .quantization import (
-    quantize_dynamic,
-    prepare_qat_fx,
-    convert_qat_fx,
+
+# 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,
 )
-from .training import train_loop
-from .utils import save_model, load_model, set_dropout
-from .hf_checkpoint import hf_login, save_checkpoint, download_checkpoint
+
+# 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__ = [
-    "PositionalEncoding",
+    # Core model components
     "BitTransformerLM",
+    "PositionalEncoding",
     "ReversibleLoggingTransformerEncoderLayer",
-    "example_usage",
+    "diffusion_inference",
     "example_training_step",
-    "TelemetrySynthesizer",
-    "detect_metric_drift",
-    "collapse_submodel",
-    "save_distilled_model",
-    "hil_safe_inference",
-    "demo_hil_safety",
-    "safe_sample_with_retry",
-    "text_to_bits",
+    "example_usage",
+    "infer_long_sequence",
+
+    # I/O and data processing
     "bits_to_text",
-    "infer_text",
-    "enforce_parity",
-    "plot_telemetry",
-    "run_dashboard",
-    "configure_optimizer",
-    "adjust_learning_rate",
-    "expand_model",
-    "distill_step",
-    "TelemetryLog",
-    "quantize_dynamic",
-    "prepare_qat_fx",
-    "convert_qat_fx",
-    "train_loop",
-    "wrap_fsdp",
-    "make_pipeline",
     "compress_bits",
     "decompress_bits",
+    "enforce_parity",
+    "infer_text",
     "model_output_decompress",
     "pack_bits",
+    "text_to_bits",
     "unpack_bits",
-    "infer_long_sequence",
-    "diffusion_inference",
-    "save_model",
-    "load_model",
-    "set_dropout",
+
+    # 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",
-    "download_checkpoint",
-    "cpu_autocast",
+    "save_model",
+    "set_dropout",
 ]