#!/usr/bin/env python3
"""Test different sampling parameters to fix repetition bug."""

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

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

# Load current model
config = TrainingConfig()
config.embed_dim = 768
config.sequence_length = 512
config.checkpoint_path = "C:\\Users\\shank\\AppData\\Local\\Programs\\Windsurf\\storage\\proto_lm\\checkpoints\\step00003652-20251021-183956\\checkpoint.pt"
config.device = "cuda"

print("Loading model...")
model = ProtoLanguageModel(config)
print(f"Loaded: step {model.step}, vocab {model.vocab.size()}")

# Test different sampling parameters
test_prompts = ["Hello", "I am", "I"]

print("\n" + "="*80)
print("TEST 1: Default parameters (temp=0.8, rep_pen=1.1, top_k=50, top_p=0.9)")
print("="*80)
for prompt in test_prompts:
    result = model.sample(prompt, max_tokens=20, temperature=0.8, repetition_penalty=1.1, top_k=50, top_p=0.9)
    print(f"{prompt:15s} → {result}")

print("\n" + "="*80)
print("TEST 2: Higher temperature (temp=1.2)")
print("="*80)
for prompt in test_prompts:
    result = model.sample(prompt, max_tokens=20, temperature=1.2, repetition_penalty=1.1, top_k=50, top_p=0.9)
    print(f"{prompt:15s} → {result}")

print("\n" + "="*80)
print("TEST 3: Stronger repetition penalty (rep_pen=1.5)")
print("="*80)
for prompt in test_prompts:
    result = model.sample(prompt, max_tokens=20, temperature=0.8, repetition_penalty=1.5, top_k=50, top_p=0.9)
    print(f"{prompt:15s} → {result}")

print("\n" + "="*80)
print("TEST 4: Lower top_k (top_k=10)")
print("="*80)
for prompt in test_prompts:
    result = model.sample(prompt, max_tokens=20, temperature=0.8, repetition_penalty=1.1, top_k=10, top_p=0.9)
    print(f"{prompt:15s} → {result}")

print("\n" + "="*80)
print("TEST 5: Combined: temp=1.2, rep_pen=2.0, top_k=20")
print("="*80)
for prompt in test_prompts:
    result = model.sample(prompt, max_tokens=20, temperature=1.2, repetition_penalty=2.0, top_k=20, top_p=0.9)
    print(f"{prompt:15s} → {result}")

print("\nDone!")