#!/usr/bin/env python3
"""Test the zero-start checkpoint after corpus training."""

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

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

print("="*80)
print("TESTING ZERO-START CHECKPOINT")
print("="*80)
print()

cfg = TrainingConfig()
cfg.checkpoint_path = None
model = ProtoLanguageModel(cfg, learning_enabled=False)
model.load_checkpoint('storage/proto_lm/zero_start.pt')

print(f"Checkpoint loaded:")
print(f"  Step: {model.step}")
print(f"  Vocab: {model.vocab.size()}")
print(f"  Merges: {len(model.vocab.merges)}")
print()

# Test with prompts matching the training data
test_prompts = [
    "Hello",
    "Hi",
    "What is 3 plus 4",
    "Thank you",
    "I am",
]

print("GENERATION TESTS")
print("="*80)
print()

for prompt in test_prompts:
    output = model.sample(
        prompt,
        max_tokens=20,
        temperature=0.6,  # Lower for more deterministic
        repetition_penalty=2.0,  # Strong penalty
        top_p=0.9,
        top_k=50,
    )
    print(f"Input:  '{prompt}'")
    print(f"Output: '{output}'")
    print()

print("="*80)