File size: 3,226 Bytes
99735a3 | 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 | #!/usr/bin/env python
"""PALIMPSESTE — Train fluid conversation model with expanded corpus.
Combines killer_corpus + conversation_corpus + trivia for maximum breadth.
"""
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
# Load ALL data
pairs = list(get_killer_corpus())
pairs.extend(list(get_conversation_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]])
print(f'Loaded {len(trivia[:3000])} trivia pairs', flush=True)
except FileNotFoundError:
pass
print(f'Total pairs: {len(pairs)}', flush=True)
# Config
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)
lm = PalimpsesteForCausalLM(config=cfg)
# BPE tokenizer
print('Training BPE tokenizer...', flush=True)
full_text = ' '.join(q + ' ' + a for q, a in pairs)
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=2000)
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(pairs, 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 — conversation, creative, code, summarization
print('\n=== FLUID CONVERSATION TESTS ===', flush=True)
tests = [
# Core Q/A
'who are you',
'what is python',
# Conversation flow
'that is interesting',
'tell me more',
'thanks for the help',
# Creative writing
'write a poem about the sea',
'write a haiku',
'describe a sunset',
# Code
'write a python function to reverse a string',
'how to read a file in python',
# Summarization
'summarize this: python is a high-level programming language known for its clear syntax and large standard library. it supports multiple paradigms including object-oriented and functional programming.',
# Opinions
'why is python popular',
'what is the best programming language',
# Emotional
'i am sad',
'you are amazing',
]
for q in tests:
t0 = time.perf_counter()
r = lm.respond(q, max_new_tokens=200)
dt = time.perf_counter() - t0
print(f' [{dt:.1f}s] {q}', flush=True)
print(f' -> {r}', 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)
print('DONE', flush=True)
|