#!/usr/bin/env python3
"""Test the NEWLY trained 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 NEWLY TRAINED CHECKPOINT (synthetic_baseline.pt)")
print("="*80)
print()

# Load the synthetic_baseline checkpoint
cfg = TrainingConfig()
cfg.checkpoint_path = 'storage/proto_lm/synthetic_baseline.pt'
model = ProtoLanguageModel(cfg, learning_enabled=False)
loaded = model.load_checkpoint('storage/proto_lm/synthetic_baseline.pt')

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

# Test prompts that match training data
test_prompts = [
    "Hello",
    "What's 3 plus 4",
    "Hi",
    "Thank you",
    "Summarize",
]

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

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

print("="*80)

# Check vocabulary
print("\nVocabulary sample:")
print(f"Total: {model.vocab.size()}")
print(f"First 50: {model.vocab.tokens[:50]}")
if model.vocab.size() > 500:
    print(f"Tokens 500-550: {model.vocab.tokens[500:550]}")