| |
| """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 |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| print('\n=== FLUID CONVERSATION TESTS ===', flush=True) |
| tests = [ |
| |
| 'who are you', |
| 'what is python', |
| |
| 'that is interesting', |
| 'tell me more', |
| 'thanks for the help', |
| |
| 'write a poem about the sea', |
| 'write a haiku', |
| 'describe a sunset', |
| |
| 'write a python function to reverse a string', |
| 'how to read a file in python', |
| |
| '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.', |
| |
| 'why is python popular', |
| 'what is the best programming language', |
| |
| '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) |
|
|
| |
| 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) |
|
|
