#!/usr/bin/env python3
"""Test why generation is broken even when model learns."""
import sys
import torch
sys.path.insert(0, 'c:/MONIKA')
from salience_os_seed.proto_lm._torch_impl import ProtoLanguageModel, TrainingConfig

m = ProtoLanguageModel(TrainingConfig(checkpoint_path=None))

# Train heavily
phrase = "Hello I am Monika"
for _ in range(100):
    m.training_step(phrase)

print(f"Loss: {m._latest_loss:.6f}")

# Test forward pass
ids = m.encode("Hello I am", mutate=False)
print(f"\nEncoded 'Hello I am': {ids}")
print(f"Tokens: {[m.vocab.tokens[i] for i in ids]}")

# Forward pass
with torch.no_grad():
    m.eval()
    context = torch.tensor(ids, device=m.device).unsqueeze(0)
    logits = m._forward_logits(context)
    next_token_logits = logits[0, -1, :]
    
    # Greedy decode
    next_id = next_token_logits.argmax().item()
    print(f"\nNext token ID: {next_id}")
    print(f"Next token: '{m.vocab.tokens[next_id]}'")
    print(f"Vocab size: {m.vocab.size()}")
    
# Try sample
result = m.sample("Hello I am", max_tokens=5, temperature=0.1, repetition_penalty=1.0)
print(f"\nSample output: '{result}'")