🚀 Refined BitTransformerLM: Organized codebase with best practices

bit_transformer/bit_io.py

@@ -1,6 +1,9 @@
-from typing import List, TYPE_CHECKING
-import torch
+"""Text-to-bit conversion utilities with parity protection."""
+
 import sys
+from typing import TYPE_CHECKING, List
+
+import torch
 
 try:  # torch.compile may be unavailable or unsupported
     if torch.__version__ and tuple(map(int, torch.__version__.split(".")[:2])) >= (2, 0) and sys.version_info < (3, 11):
@@ -21,7 +24,14 @@ if TYPE_CHECKING:  # pragma: no cover
 
 @compile_fn
 def bytes_to_bits(data: bytes) -> List[int]:
-    """Convert bytes to bits with per-byte parity bit."""
+    """Convert bytes to bits with per-byte parity bit.
+
+    Args:
+        data: Input bytes to convert.
+
+    Returns:
+        List of bits with parity protection (9 bits per byte).
+    """
     result: List[int] = []
     for b in data:
         bits = [(b >> i) & 1 for i in reversed(range(8))]
@@ -32,7 +42,17 @@ def bytes_to_bits(data: bytes) -> List[int]:
 
 @compile_fn
 def bits_to_bytes(bits: List[int]) -> bytes:
-    """Convert parity-protected bits back to bytes."""
+    """Convert parity-protected bits back to bytes.
+
+    Args:
+        bits: List of bits with parity protection (length must be multiple of 9).
+
+    Returns:
+        Decoded bytes.
+
+    Raises:
+        ValueError: If bit stream length is not multiple of 9 or parity check fails.
+    """
     if len(bits) % 9 != 0:
         raise ValueError("Bit stream length must be multiple of 9")
     out = bytearray()
@@ -50,10 +70,26 @@ def bits_to_bytes(bits: List[int]) -> bytes:
 
 
 def text_to_bits(text: str) -> List[int]:
+    """Convert text to parity-protected bit sequence.
+
+    Args:
+        text: Input text string.
+
+    Returns:
+        List of bits with parity protection.
+    """
     return bytes_to_bits(text.encode("utf-8"))
 
 
 def bits_to_text(bits: List[int]) -> str:
+    """Convert parity-protected bit sequence to text.
+
+    Args:
+        bits: List of bits with parity protection.
+
+    Returns:
+        Decoded text string (with error replacement for invalid UTF-8).
+    """
     return bits_to_bytes(bits).decode("utf-8", errors="replace")