#!/usr/bin/env python3
import sys
sys.path.insert(0, 'C:\\MONIKA')

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

cfg = TrainingConfig()
cfg.checkpoint_path = 'storage/proto_lm/checkpoint.pt'
m = ProtoLanguageModel(cfg)
loaded = m.load_checkpoint()

print(f"Checkpoint loaded: {loaded}")
print(f"Step: {m.step}")
print(f"Vocab size: {m.vocab.size()}")
print(f"Merges: {len(m.vocab.merges)}")
print(f"First 30 tokens: {m.vocab.tokens[:30]}")
print(f"Sample from vocab[200:220]: {m.vocab.tokens[200:220]}")

# Test generation
output = m.sample("Hello I am", max_tokens=15, temperature=0.8)
print(f"\nGeneration test:")
print(f"Input: 'Hello I am'")
print(f"Output: '{output}'")