#!/usr/bin/env python3
"""Train with new word-level vocab and save checkpoint for MCP."""
import sys
import torch
sys.path.insert(0, 'c:/MONIKA')
from salience_os_seed.proto_lm._torch_impl import ProtoLanguageModel, TrainingConfig

# Create model with new vocab
m = ProtoLanguageModel(TrainingConfig(checkpoint_path=None))

print(f"Vocab size: {m.vocab.size()}")
print(f"Sample tokens: {m.vocab.tokens[116:126]}")

# Train on diverse conversational data
training_data = [
    "Hello I am Monika",
    "Hello I am learning to speak",
    "I am learning to communicate",
    "I want to learn more",
    "I can help you",
    "Thank you very much",
    "I am here to help",
    "I will try my best",
    "I think I understand",
    "I feel happy today",
    "You are very kind",
    "This is very good",
    "I want to know more",
    "Tell me about that",
    "I am grateful for this",
] * 20

for text in training_data:
    m.training_step(text)

print(f"\nTraining loss: {m._latest_loss:.4f}")
print(f"Step: {m.step}")

# Test generation
print("\nGeneration tests:")
for prompt in ["Hello I am", "I am", "I want to"]:
    result = m.sample(prompt, max_tokens=5, temperature=0.3, repetition_penalty=1.3)
    print(f"  '{prompt}' -> '{result}'")

# Save checkpoint
torch.save({
    'model': m.state_dict(),
    'optimizer': m.optimizer.state_dict(),
    'step': m.step,
    'vocab': {'tokens': m.vocab.tokens, 'merges': m.vocab.merges},
    'scheduler': None
}, 'storage/monika_coherent.pt')

print("\n✓ Saved to storage/monika_coherent.pt")