Remove nested directory: BitTransformerLM/bit_transformer/bit_io.py

BitTransformerLM/bit_transformer/bit_io.py

@@ -1,97 +0,0 @@
-from typing import List, TYPE_CHECKING
-import torch
-import sys
-
-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):
-        compile_fn = torch.compile
-    else:
-        raise RuntimeError
-except Exception:  # pragma: no cover
-
-    def compile_fn(fn=None, **kwargs):
-        if fn is None:
-            return lambda f: f
-        return fn
-
-
-if TYPE_CHECKING:  # pragma: no cover
-    from .model import BitTransformerLM
-
-
-@compile_fn
-def bytes_to_bits(data: bytes) -> List[int]:
-    """Convert bytes to bits with per-byte parity bit."""
-    result: List[int] = []
-    for b in data:
-        bits = [(b >> i) & 1 for i in reversed(range(8))]
-        parity = sum(bits) % 2
-        result.extend(bits + [parity])
-    return result
-
-
-@compile_fn
-def bits_to_bytes(bits: List[int]) -> bytes:
-    """Convert parity-protected bits back to bytes."""
-    if len(bits) % 9 != 0:
-        raise ValueError("Bit stream length must be multiple of 9")
-    out = bytearray()
-    for i in range(0, len(bits), 9):
-        chunk = bits[i : i + 9]
-        payload = chunk[:8]
-        parity = chunk[8]
-        if parity != sum(payload) % 2:
-            raise ValueError("Parity check failed")
-        value = 0
-        for bit in payload:
-            value = (value << 1) | bit
-        out.append(value)
-    return bytes(out)
-
-
-def text_to_bits(text: str) -> List[int]:
-    return bytes_to_bits(text.encode("utf-8"))
-
-
-def bits_to_text(bits: List[int]) -> str:
-    return bits_to_bytes(bits).decode("utf-8", errors="replace")
-
-
-def infer_text(
-    model: "BitTransformerLM",
-    text: str,
-    c_floor: float = 0.3,
-    s_floor: float = 0.5,
-) -> str:
-    """Run text through the model using the safety gate."""
-    from .safety import hil_safe_inference
-    bits = text_to_bits(text)
-    tensor = torch.tensor(bits, dtype=torch.long).unsqueeze(0)
-    out_bits, _ = hil_safe_inference(model, tensor, c_floor=c_floor, s_floor=s_floor)
-    return bits_to_text(out_bits.squeeze(0).tolist())
-
-
-def sample_text(
-    model: "BitTransformerLM",
-    prompt: str,
-    max_new_tokens: int = 16,
-    temperature: float = 1.0,
-    top_p: float = 1.0,
-) -> str:
-    """Generate text from the model using simple top-p sampling."""
-    model.eval()
-    bits = text_to_bits(prompt)
-    tensor = torch.tensor(bits, dtype=torch.long).unsqueeze(0)
-    for _ in range(max_new_tokens * 9):
-        if tensor.size(1) >= model.pos_enc.pe.size(0):
-            break
-        logits, _ = model(tensor, causal=True)
-        prob = logits[0, -1].softmax(-1) / temperature
-        sorted_prob, sorted_idx = prob.sort(descending=True)
-        cumulative = sorted_prob.cumsum(0)
-        mask = cumulative > top_p
-        sorted_prob[mask] = 0
-        sorted_prob = sorted_prob / sorted_prob.sum()
-        next_bit = sorted_idx[torch.multinomial(sorted_prob, 1)]
-        tensor = torch.cat([tensor, next_bit.view(1, 1)], dim=1)
-    return bits_to_text(tensor.squeeze(0).tolist())