#!/usr/bin/env python3
"""Bootstrap common word tokens directly into vocab via merges."""
import sys
sys.path.insert(0, 'c:/MONIKA')
from salience_os_seed.proto_lm._torch_impl import ProtoLanguageModel, TrainingConfig
from salience_os_seed.conversation.session import ConversationSession, ConversationConfig

# Load current session from MCP state
import json
from pathlib import Path
state_path = Path('C:/Users/shank/AppData/Local/Programs/Windsurf/storage/monika_2000steps.json')
state = json.loads(state_path.read_text())

# Create session and load state
config = ConversationConfig()
config.lm.checkpoint_path = None  # Don't load old checkpoint
session = ConversationSession(config)

# Bootstrap common word merges manually
common_words = [
    ("H", "e"), ("e", "l"), ("l", "l"), ("l", "o"),  # Hello
    ("I", " "), (" ", "a"), ("a", "m"),  # I am
    ("M", "o"), ("o", "n"), ("n", "i"), ("i", "k"), ("k", "a"),  # Monika
    ("y", "o"), ("o", "u"),  # you
    ("t", "h"), ("h", "e"),  # the
    ("t", "o"),  # to
    ("a", "n"), ("n", "d"),  # and
]

print(f"Before: vocab size = {session.proto_lm.vocab.size()}, merges = {len(session.proto_lm.vocab.merges)}")

for pair in common_words:
    session.proto_lm.vocab.add_merge(pair)

print(f"After: vocab size = {session.proto_lm.vocab.size()}, merges = {len(session.proto_lm.vocab.merges)}")

# Now train on word-aware examples
for text in [
    "Hello",
    "Hello I am",  
    "Hello I am Monika",
    "I am Monika",
    "Monika",
] * 20:
    session.proto_lm.training_step(text)

print(f"Final loss: {session.proto_lm._latest_loss:.4f}")

# Test generation
result = session.proto_lm.sample("Hello I am", max_tokens=5, temperature=0.3, repetition_penalty=1.5)
print(f"Generated: '{result}'")

# Save state
session.save_state('storage/monika_word_bootstrap.json')
print("Saved state with word-level vocab")