#!/usr/bin/env python3
"""Initialize and train a larger MONIKA model with proper scaling."""

import sys
sys.path.insert(0, 'C:\\MONIKA')

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

# Scale up config for RTX 5060 16GB
config = TrainingConfig()
config.embed_dim = 384  # 3x larger (128 -> 384)
config.sequence_length = 256  # 2x longer context
config.checkpoint_path = "storage/proto_lm/large_model.pt"
config.device = "cuda"
config.learning_rate = 3e-4  # Adjust for larger model

print("=== Initializing Larger MONIKA Model ===")
print(f"Embed dim: {config.embed_dim}")
print(f"Sequence length: {config.sequence_length}")
print(f"Device: {config.device}")

model = ProtoLanguageModel(config)

print(f"\nInitial state:")
print(f"  Vocab size: {model.vocab.size()}")
print(f"  Embed capacity: {model.embed.num_embeddings}")
print(f"  Embed dim: {model.embed.weight.shape[1]}")
print(f"  Output capacity: {model.output.out_features}")

# Count parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"  Total parameters: {total_params:,}")

# Test training step
print("\nTesting training...")
loss = model.training_step("The quick brown fox jumps over the lazy dog.")
print(f"  First training step loss: {loss:.4f}")

# Save initial checkpoint
checkpoint_path = model.save_checkpoint(reason="large-model-init", tags=["larger-architecture", "384-dim"])
print(f"\nSaved checkpoint to: {checkpoint_path}")

print("\n=== Model ready for training ===")
print("You can now use MCP tools to train this larger model.")