🚀 Refined BitTransformerLM: Organized codebase with best practices

scripts/testing/code_test.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""
+Test BitTransformerLM on Code/Math Completion
+"""
+
+import sys
+import torch
+import torch.nn.functional as F
+
+sys.path.append('/data')
+sys.path.append('/data/BitTransformerLM')
+
+from bit_transformer import BitTransformerLM, text_to_bits, bits_to_text
+
+def load_model():
+    model = BitTransformerLM(
+        d_model=512, nhead=16, num_layers=8, dim_feedforward=1024,
+        max_seq_len=512, reversible=True, use_checkpoint=False,
+        use_autocast=False, use_act=True, act_threshold=0.9,
+        lambda_K=0.05, lambda_C=0.05, lambda_S=0.05
+    )
+    
+    checkpoint = torch.load('/data/BitTransformerLM/checkpoints/checkpoint_best.pt', map_location='cpu')
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    
+    return model
+
+def code_generate(model, prompt, max_chars=10):
+    """Generate code/math completion."""
+    print(f"\n🧮 Code completion for: '{prompt}'")
+    
+    input_bits = text_to_bits(prompt)
+    generated_bits = input_bits.copy()
+    
+    results = []
+    
+    with torch.no_grad():
+        for char_idx in range(max_chars):
+            # Generate 9 bits for one character
+            char_bits = []
+            
+            for bit_idx in range(9):
+                context = generated_bits + char_bits
+                context = context[-400:] if len(context) > 400 else context
+                context_tensor = torch.tensor(context, dtype=torch.long).unsqueeze(0)
+                
+                logits, telemetry = model(context_tensor)
+                next_bit_logits = logits[0, -1, :]
+                
+                if bit_idx < 8:  # Data bits
+                    # Use different sampling for code (more deterministic)
+                    temperature = 0.5  # Lower temperature for code
+                    next_bit_logits = next_bit_logits / temperature
+                    
+                    # Greedy sampling for first few characters to see most likely
+                    if char_idx < 3:
+                        next_bit = torch.argmax(next_bit_logits).item()
+                    else:
+                        probs = F.softmax(next_bit_logits, dim=-1)
+                        next_bit = torch.multinomial(probs, 1).item()
+                else:  # Parity bit
+                    data_bits = char_bits[:8]
+                    expected_parity = sum(data_bits) % 2
+                    next_bit = expected_parity
+                
+                char_bits.append(next_bit)
+            
+            # Add character and try to decode
+            generated_bits.extend(char_bits)
+            
+            # Decode this character
+            data_bits = char_bits[:8]
+            byte_val = sum(bit * (2**(7-i)) for i, bit in enumerate(data_bits))
+            
+            if 32 <= byte_val <= 126:  # Printable ASCII
+                char = chr(byte_val)
+                print(f"  +'{char}' (confidence: {torch.max(F.softmax(next_bit_logits, dim=-1)).item():.3f})")
+                results.append(char)
+                
+                # Stop on natural code endings
+                if char in ';{}()[]':
+                    break
+            else:
+                print(f"  +[{byte_val}] (non-printable)")
+                results.append('?')
+    
+    completion = ''.join(results)
+    print(f"✨ Result: '{prompt}' → '{prompt}{completion}'")
+    
+    return completion
+
+def main():
+    print("🚀 BITRANSFORMERLM CODE/MATH COMPLETION TEST")
+    print("=" * 50)
+    
+    model = load_model()
+    print("✅ Model loaded!")
+    
+    # Test structured prompts that might have learned patterns
+    test_cases = [
+        # Math equations
+        "2 + 2 =",
+        "1 + 1 =", 
+        "5 * 3 =",
+        "10 / 2 =",
+        
+        # Simple code patterns
+        "def hello():",
+        "if x ==",
+        "for i in",
+        "print(",
+        "return",
+        
+        # Simple patterns
+        "a, b, c,", 
+        "1, 2, 3,",
+        "red, blue,",
+        
+        # HTML/markup
+        "<div>",
+        "function(",
+        "var x =",
+    ]
+    
+    print(f"\n🧮 Testing {len(test_cases)} code/math patterns:")
+    
+    for i, prompt in enumerate(test_cases):
+        print(f"\n--- Test {i+1}/{len(test_cases)} ---")
+        completion = code_generate(model, prompt, max_chars=6)
+        
+        # Quick analysis
+        if any(c.isalnum() for c in completion):
+            print("   📝 Contains alphanumeric - GOOD!")
+        if any(c in "0123456789" for c in completion):
+            print("   🔢 Contains numbers - EXCELLENT!")
+        if any(c in "=(){}[];," for c in completion):
+            print("   💻 Contains code symbols - PROMISING!")
+
+if __name__ == "__main__":
+    main()