File size: 4,686 Bytes
9202d01 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 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() |