File size: 3,988 Bytes
a61306f | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | #!/usr/bin/env python
"""PALIMPSESTE — Train 100K model. Target: ~100K pairs, ~4-5M tokens, ~2-3GB."""
import json, sys, time, os, random
sys.path.insert(0, '.')
sys.path.insert(0, 'examples')
random.seed(42)
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.hf import HFPalimpsesteLM
from palimseste.bpe import BPETokenizer
from generate_100k import (
load_existing, gen_definitions, gen_geography_expanded,
gen_code_expanded, gen_conversation_variants, gen_tech_definitions,
gen_history_facts,
)
# Build corpus — controlled size
pairs = []
pairs.extend(load_existing())
pairs.extend(gen_definitions())
pairs.extend(gen_geography_expanded())
pairs.extend(gen_code_expanded())
pairs.extend(gen_conversation_variants())
pairs.extend(gen_tech_definitions())
pairs.extend(gen_history_facts())
# Add controlled math tables (enough to reach 100K without bloat)
for i in range(1, 50):
for j in range(1, 50):
pairs.append((f'what is {i} times {j}', f'{i} times {j} is {i*j}'))
for i in range(1, 100):
for j in range(1, 100):
pairs.append((f'what is {i} plus {j}', f'{i} plus {j} is {i+j}'))
for i in range(1, 30):
pairs.append((f'what is {i} squared', f'{i} squared is {i*i}'))
if i <= 12:
pairs.append((f'what is {i} cubed', f'{i} cubed is {i**3}'))
# Add augmented variants (trailing space)
augmented = []
for q, a in pairs:
augmented.append((q, a))
augmented.append((q + ' ', a))
# Deduplicate
seen = set()
final = []
for q, a in augmented:
k = q.lower().strip()
if k not in seen:
seen.add(k)
final.append((q, a))
# Cap at 100K
if len(final) > 100000:
final = final[:100000]
total_chars = sum(len(q) + len(a) for q, a in final)
print(f'Corpus: {len(final):,} pairs, {total_chars:,} chars', flush=True)
print(f'Estimated tokens: ~{total_chars // 2:,}', flush=True)
# Free disk space
if os.path.exists('palimpseste-max/palimpseste_memory.bin'):
os.remove('palimpseste-max/palimpseste_memory.bin')
print('Deleted old model', flush=True)
# Config
cfg = PalimpsesteConfig(D=20_000, context_window=256, kernel_radius=400, temperature=0.0)
lm = PalimpsesteForCausalLM(config=cfg)
# BPE
print('Training BPE...', flush=True)
full_text = ' '.join(q + ' ' + a for q, a in final)
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=3000)
bpe.train(full_text, verbose=False)
lm.attach_tokenizer(bpe)
print(f'BPE vocab: {bpe.vocab_size_actual}', flush=True)
# Train
print('Training...', flush=True)
t0 = time.perf_counter()
n = lm.train_on_qa_pairs(final, verbose=False)
dt = time.perf_counter() - t0
print(f'Trained: {n:,} tokens in {dt:.0f}s ({n/dt:.0f} tok/s)', flush=True)
print(f'|M| = {len(lm.mem):,}', flush=True)
# Test
print('\n=== TESTS ===', flush=True)
for q in ['who are you', 'what is python', 'what is the capital of japan', 'what is 25 times 13', 'who was einstein']:
r = lm.respond(q, max_new_tokens=200)
print(f' {q} -> {r[:60]}', flush=True)
# Save
print('\nSaving...', 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)
sz = os.path.getsize('./palimpseste-max/palimpseste_memory.bin')
print(f'Saved: {sz/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
hf.mem = lm.mem
hf.save_pretrained('./palimpseste-max', tune_lsh=False)
print(f'Final: {os.path.getsize("./palimpseste-max/palimpseste_memory.bin")/1024/1024:.0f} MB', flush=True)
# Speed
print('\n=== SPEED ===', flush=True)
for q in ['who are you', 'hello']:
t0 = time.perf_counter()
r = lm.respond(q, max_new_tokens=100)
print(f' [{time.perf_counter()-t0:.1f}s] {q} -> {r[:50]}', flush=True)
print('DONE', flush=True)
|