File size: 2,497 Bytes
8073a30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | #!/usr/bin/env python
"""PALIMPSESTE — Mega training script.
Trains with context_window=256, D=20000, ALL available data, temp=0.0.
This is the training that produces the killer model.
"""
import json, sys, time
sys.path.insert(0, '.')
sys.path.insert(0, 'examples')
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.hf import HFPalimpsesteLM
from corpus_chat import get_corpus
# Load ALL data
pairs = list(get_corpus())
with open('trivia_qa_pairs.json') as f:
pairs.extend([(item['q'], item['a']) for item in json.load(f)])
with open('large_dataset.json') as f:
pairs.extend([(item['q'], item['a']) for item in json.load(f)])
print(f'Total pairs: {len(pairs)}', flush=True)
# Config — the game changer
cfg = PalimpsesteConfig(
D=20_000,
context_window=256,
kernel_radius=400,
kernel_min_weight=1e-6,
temperature=0.0,
)
print(f'Config: D={cfg.D} ctx={cfg.context_window} temp={cfg.temperature}', flush=True)
# Build model
lm = PalimpsesteForCausalLM(config=cfg)
full_text = ''.join(q + a for q, a in pairs)
lm.build_tokenizer(full_text)
print(f'Vocab: {lm.tokenizer.vocab_size}', flush=True)
# TRAIN
print('Training...', flush=True)
t0 = time.perf_counter()
n_tokens = lm.train_on_qa_pairs(pairs, verbose=True)
dt = time.perf_counter() - t0
print(f'Trained: {n_tokens:,} tokens in {dt:.1f}s ({n_tokens/dt:.0f} tok/s)', flush=True)
print(f'|M| = {len(lm.mem):,} traces', flush=True)
# TEST
print('\n=== TESTING RETRAINED MODEL ===', flush=True)
tests = [
'who are you', 'hello', 'what is python', 'how do you learn',
'do you use a gpu', 'what is the capital of france',
'what can you do', 'are you conscious',
'who won super bowl xx', 'what is a black hole',
'what is recursion', 'are you alive',
]
for q in tests:
resp = lm.respond(q, max_new_tokens=200)
print(f' Q: {q}', flush=True)
print(f' A: {resp}', flush=True)
print(flush=True)
# SAVE
print('Saving model...', flush=True)
# HFPalimpsesteLM fields are init=False, so we create from config then copy state
hf_lm = HFPalimpsesteLM(config=lm.config)
hf_lm.mem = lm.mem
hf_lm.phi = lm.phi
hf_lm.encoder = lm.encoder
hf_lm.tokenizer = lm.tokenizer
hf_lm._self_hv = lm._self_hv
hf_lm._token_bits_cache = None
hf_lm.save_pretrained('./palimpseste-max', tune_lsh=True)
import os
size_mb = os.path.getsize('./palimpseste-max/palimpseste_memory.bin') / 1024 / 1024
print(f'Saved: {size_mb:.0f} MB', flush=True)
print('DONE', flush=True)
|