#!/usr/bin/env python3
"""Bulk train MONIKA on entire corpus."""

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

from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig
from pathlib import Path

# Load model
config = TrainingConfig()
config.embed_dim = 768
config.sequence_length = 512
config.checkpoint_path = "storage/proto_lm/massive_model.pt"
config.device = "cuda"

print("Loading model...")
model = ProtoLanguageModel(config)
model.load_checkpoint("storage/proto_lm/checkpoints/step00000001-20251021-111117/checkpoint.pt")

print(f"Model loaded: step {model.step}, vocab {model.vocab.size()}")

# Load corpus
corpus_path = Path("C:\\MONIKA\\training_corpus.txt")
corpus = corpus_path.read_text(encoding='utf-8')
sentences = [s.strip() for s in corpus.split('\n') if s.strip()]

print(f"\nCorpus: {len(sentences)} sentences, {len(corpus)} chars")

# Train on entire corpus
print("\nTraining...")
for i, sentence in enumerate(sentences):
    loss = model.training_step(sentence)
    if (i + 1) % 20 == 0:
        print(f"  Step {model.step}: loss {loss:.4f}, vocab {model.vocab.size()}")

print(f"\nFinal: step {model.step}, vocab {model.vocab.size()}")

# Test generation
print("\nTesting generation...")
result = model.sample("I am learning to", max_tokens=50)
print(f"Generated: {result[:200]}")

# Save
checkpoint = model.save_checkpoint(reason="bulk-trained-on-corpus", tags=["corpus-trained", "768-dim"])
print(f"\nSaved: {checkpoint}")