#!/usr/bin/env python
"""PALIMPSESTE — Train massive model with all corpora combined."""
import json, sys, time, os
sys.path.insert(0, '.')
sys.path.insert(0, 'examples')

from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.hf import HFPalimpsesteLM
from palimseste.bpe import BPETokenizer
from killer_corpus import get_killer_corpus
from conversation_corpus import get_conversation_corpus
from massive_corpus import get_massive_corpus

pairs = list(get_killer_corpus())
pairs.extend(list(get_conversation_corpus()))
pairs.extend(list(get_massive_corpus()))

try:
    with open('trivia_qa_pairs.json') as f:
        trivia = json.load(f)
    pairs.extend([(item['q'].lower(), item['a'].lower()) for item in trivia[:3000]])
except FileNotFoundError:
    pass

print(f'Total pairs: {len(pairs)}', flush=True)

cfg = PalimpsesteConfig(D=20_000, context_window=256, kernel_radius=400, temperature=0.0)
lm = PalimpsesteForCausalLM(config=cfg)

full_text = ' '.join(q + ' ' + a for q, a in pairs)
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=2500)
bpe.train(full_text, verbose=False)
lm.attach_tokenizer(bpe)
print(f'BPE vocab: {bpe.vocab_size_actual}', flush=True)

print('Training...', flush=True)
t0 = time.perf_counter()
n = lm.train_on_qa_pairs(pairs, verbose=False)
print(f'Trained: {n:,} tokens in {time.perf_counter()-t0:.0f}s, |M|={len(lm.mem):,}', flush=True)

# Quick test
for q in ['who are you', 'write a poem about the moon', 'how to stay focused', 'what is dark matter']:
    r = lm.respond(q, max_new_tokens=200)
    print(f'  {q} -> {r[:60]}', flush=True)

# Save
print('Saving...', flush=True)
hf = HFPalimpsesteLM(config=lm.config)
hf.mem = lm.mem; hf.phi = lm.phi; hf.encoder = lm.encoder
hf.tokenizer = lm.tokenizer; hf._self_hv = lm._self_hv; hf._token_bits_cache = None
hf.save_pretrained('./palimpseste-max', tune_lsh=False)
print(f'Saved: {os.path.getsize("./palimpseste-max/palimpseste_memory.bin")/1024/1024:.0f} MB', flush=True)

# Tune LSH
print('Tuning LSH...', flush=True)
t0 = time.perf_counter()
lm.tune_lsh()
print(f'LSH tuned in {time.perf_counter()-t0:.0f}s', flush=True)

# Re-save with tuned LSH
lm.save_pretrained('./palimpseste-max', tune_lsh=False)
print(f'Final: {os.path.getsize("./palimpseste-max/palimpseste_memory.bin")/1024/1024:.0f} MB', flush=True)
print('DONE', flush=True)