File size: 11,681 Bytes
ba1fad5 |
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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 |
#!/usr/bin/env python3
"""
Full Attention BitTransformerLM Diffusion Inference Test
========================================================
Test the newly trained full bi-directional attention BitTransformerLM model
using denoising diffusion generation to evaluate improvements from full attention training.
Model Configuration:
- Same full bi-directional unchunked attention as training (chunk_size=None)
- Proper eval() mode with dropout management
- Use latest checkpoint_best.pt from full attention training
- Test with same diffusion inference that worked before
"""
import sys
import torch
import torch.nn.functional as F
from datetime import datetime
sys.path.append('/data')
sys.path.append('/data/BitTransformerLM')
from bit_transformer import (
BitTransformerLM,
text_to_bits,
bits_to_text,
diffusion_inference,
set_dropout
)
def load_full_attention_model():
"""Load the newly trained full attention BitTransformerLM model."""
print("🚀 Loading Full Attention BitTransformerLM for diffusion inference...")
# Create model with SAME configuration as full attention training
model = BitTransformerLM(
d_model=512, # Same as training
nhead=16, # Same as training
num_layers=8, # Same as training
dim_feedforward=1024, # Same as training
max_seq_len=512, # Same as training
reversible=True, # Same as training
use_checkpoint=False, # Disable for inference
use_autocast=False, # Disable for inference
use_act=True, # Same as training
act_threshold=0.9, # Same as training
lambda_K=0.05, # Same as training
lambda_C=0.05, # Same as training
lambda_S=0.05, # Same as training
chunk_size=None, # FULL ATTENTION - same as training
overlap=0, # Same as training
full_attn_logging=True # Same as training
)
# Load the latest checkpoint_best.pt (should be from full attention training)
checkpoint_path = '/data/BitTransformerLM/checkpoints/checkpoint_best.pt'
checkpoint = torch.load(checkpoint_path, map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
# Set to evaluation mode with proper dropout
model.eval()
set_dropout(model, 0.0) # Disable dropout for inference
# Get checkpoint info
epoch = checkpoint.get('epoch', 'unknown')
loss = checkpoint.get('loss', 'unknown')
print(f"✅ Full Attention Model loaded! Epoch: {epoch}, Loss: {loss}")
# Calculate parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"📊 Parameters: {total_params:,}")
return model
def test_basic_diffusion_generation(model):
"""Test basic unconditional diffusion generation."""
print("\n🧪 === BASIC FULL ATTENTION DIFFUSION GENERATION ===")
results = []
test_configs = [
{"length": 36, "steps": 8, "schedule": "linear"},
{"length": 45, "steps": 12, "schedule": "cosine"},
{"length": 54, "steps": 16, "schedule": "exp"}
]
for i, config in enumerate(test_configs, 1):
print(f"\n--- Test {i}: {config['length']//9} chars, {config['schedule']} ---")
try:
# Generate with diffusion
generated_bits = diffusion_inference(
model,
length=config['length'],
steps=config['steps'],
batch_size=1,
schedule=config['schedule']
)
# Try to decode
bit_list = generated_bits.squeeze().tolist()
decoded_text = bits_to_text(bit_list)
print(f"✅ SUCCESS: '{decoded_text}'")
results.append({
"test": f"basic_{i}",
"config": config,
"success": True,
"output": decoded_text,
"bits": len(bit_list)
})
except Exception as e:
print(f"❌ FAILED: {e}")
results.append({
"test": f"basic_{i}",
"config": config,
"success": False,
"error": str(e)
})
return results
def test_conditioned_diffusion_generation(model):
"""Test prompt-conditioned diffusion generation."""
print("\n🎯 === CONDITIONED FULL ATTENTION DIFFUSION GENERATION ===")
results = []
test_prompts = [
"Hello",
"Hi there",
"What is your name?",
"The weather is",
"I am",
"Yes",
"No"
]
for prompt in test_prompts:
print(f"\n--- Prompt: '{prompt}' ---")
try:
# Convert prompt to bits
prompt_bits = text_to_bits(prompt)
# Generate continuation with diffusion (no init_bits - let it generate freely)
continuation_length = 45 # 5 character continuation
generated_bits = diffusion_inference(
model,
length=continuation_length,
steps=12,
batch_size=1,
init_bits=None,
schedule="cosine"
)
# Combine prompt + generated continuation
full_bits = prompt_bits + generated_bits.squeeze().tolist()
# Decode continuation only
continuation_bits = generated_bits.squeeze().tolist()
continuation_text = bits_to_text(continuation_bits)
# Show combined result
combined_text = prompt + continuation_text
print(f"✅ SUCCESS: '{prompt}' → '{combined_text}'")
results.append({
"test": "conditioned",
"prompt": prompt,
"success": True,
"full_output": combined_text,
"continuation": continuation_text,
"bits": len(continuation_bits)
})
except Exception as e:
print(f"❌ FAILED: {e}")
results.append({
"test": "conditioned",
"prompt": prompt,
"success": False,
"error": str(e)
})
return results
def test_code_diffusion_completion(model):
"""Test code/math completion with diffusion."""
print("\n💻 === CODE COMPLETION FULL ATTENTION DIFFUSION ===")
results = []
test_cases = [
# Math equations
"2 + 2 =",
"1 + 1 =",
"5 * 3 =",
"10 / 2 =",
# Programming constructs
"def hello():",
"if x ==",
"for i in",
"print(",
"return",
# Patterns
"a, b, c,",
"1, 2, 3,",
"function(",
"var x =",
]
for code in test_cases:
print(f"\n--- Code: '{code}' ---")
try:
# Convert to bits
code_bits = text_to_bits(code)
# Generate completion with diffusion (no init_bits)
completion_length = 45 # 5 character completion
generated_bits = diffusion_inference(
model,
length=completion_length,
steps=10,
batch_size=1,
init_bits=None,
schedule="linear"
)
# Decode completion
completion_bits = generated_bits.squeeze().tolist()
completion = bits_to_text(completion_bits)
# Show combined result
combined_text = code + completion
print(f"✅ SUCCESS: '{code}' → '{combined_text}'")
# Analyze completion
analysis = []
if any(c.isalnum() for c in completion):
analysis.append("Contains alphanumeric")
print(f" 📊 Analysis: Contains alphanumeric")
if any(c in "0123456789" for c in completion):
analysis.append("Contains numbers")
print(f" 🔢 Analysis: Contains numbers")
if any(c in "=(){}[];," for c in completion):
analysis.append("Contains code symbols")
print(f" 💻 Analysis: Contains code symbols")
results.append({
"test": "code_completion",
"prompt": code,
"success": True,
"full_output": combined_text,
"completion": completion,
"analysis": analysis,
"bits": len(completion_bits)
})
except Exception as e:
print(f"❌ FAILED: {e}")
results.append({
"test": "code_completion",
"prompt": code,
"success": False,
"error": str(e)
})
return results
def compare_with_previous_results():
"""Note about comparison with previous results."""
print("\n⚖️ === COMPARISON WITH PREVIOUS RESULTS ===")
print("Previous chunked attention model achieved:")
print("- Basic generation: 3/3 success (100%)")
print("- Conditioned generation: 7/7 success (100%)")
print("- Code completion: 13/13 success (100%)")
print("- All diffusion inference succeeded vs 0% autoregressive")
print("\nTesting if full attention training improved quality...")
def main():
print("🚀 FULL ATTENTION BITRANSFORMERLM DIFFUSION INFERENCE TEST")
print("=" * 70)
print("Testing newly trained full bi-directional attention model")
print("with denoising diffusion generation")
print("=" * 70)
# Load model
model = load_full_attention_model()
# Run tests
basic_results = test_basic_diffusion_generation(model)
conditioned_results = test_conditioned_diffusion_generation(model)
code_results = test_code_diffusion_completion(model)
# Show comparison
compare_with_previous_results()
# Calculate summary stats
total_tests = len(basic_results) + len(conditioned_results) + len(code_results)
successful_tests = sum(1 for r in basic_results + conditioned_results + code_results if r.get('success', False))
success_rate = (successful_tests / total_tests) * 100 if total_tests > 0 else 0
print(f"\n🎯 === FINAL SUMMARY ===")
print(f"Total tests: {total_tests}")
print(f"Successful: {successful_tests}")
print(f"Success rate: {success_rate:.1f}%")
print(f"\nBreakdown:")
print(f"- Basic generation: {sum(1 for r in basic_results if r.get('success', False))}/{len(basic_results)}")
print(f"- Conditioned generation: {sum(1 for r in conditioned_results if r.get('success', False))}/{len(conditioned_results)}")
print(f"- Code completion: {sum(1 for r in code_results if r.get('success', False))}/{len(code_results)}")
# Return all results for documentation
return {
'basic_results': basic_results,
'conditioned_results': conditioned_results,
'code_results': code_results,
'summary': {
'total_tests': total_tests,
'successful_tests': successful_tests,
'success_rate': success_rate,
'timestamp': datetime.now().isoformat()
}
}
if __name__ == "__main__":
results = main() |