#!/usr/bin/env python3
"""Test the synthetic baseline checkpoint."""

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

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

print("="*80)
print("TESTING SYNTHETIC BASELINE CHECKPOINT")
print("="*80)
print()

# Load checkpoint
cfg = TrainingConfig()
cfg.checkpoint_path = 'storage/proto_lm/synthetic_baseline.pt'
model = ProtoLanguageModel(cfg)
model.load_checkpoint()

print(f"Loaded checkpoint:")
print(f"  Step: {model.step}")
print(f"  Vocab: {model.vocab.size()}")
print(f"  Device: {model.device}")
print()

# Test prompts
test_prompts = [
    "Hello",
    "What is 3 plus 4",
    "Hello there",
    "Thank you",
    "I am MONIKA",
]

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

for prompt in test_prompts:
    output = model.sample(
        prompt,
        max_tokens=30,
        temperature=0.8,
        repetition_penalty=1.5,
        top_p=0.9,
        top_k=40,
    )
    print(f"Input:  '{prompt}'")
    print(f"Output: '{output}'")
    print()

print("="*80)