File size: 6,001 Bytes
5e0537e | 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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | #!/usr/bin/env python
"""PALIMPSESTE — Train a causal LM and save it (HF-compatible format).
Usage:
# train on a text file with the 'small' preset (D=10000)
python examples/train_lm.py --input corpus.txt --preset small --output ./my_model
# train on inline text
python examples/train_lm.py --text "your training text here" --preset tiny --output ./my_model
# train at 1B-scale capacity (D=100000) — large memory, slow but no GPU
python examples/train_lm.py --input corpus.txt --preset 1b --output ./palimpseste-1b
# train and push to Hugging Face Hub
python examples/train_lm.py --input corpus.txt --preset small --output ./my_model --push-to-hub user/palimpseste-small
# custom config
python examples/train_lm.py --input corpus.txt --output ./my_model --D 20000 --context-window 48 --kernel-radius 300
Training is O(1) per token — no epochs, no gradient, no GPU. The "model" is an
append-only memory of (context_hv -> token_hv) associations. The millionth
token costs exactly as much as the first.
"""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
import numpy as np
# allow running from the repo root without installing
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from palimseste.lm import PalimpsesteConfig, PRESETS
from palimseste.hf import HFPalimpsesteLM
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(
description="Train a PALIMPSESTE causal LM (no GPU, no gradient).",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=__doc__,
)
g = p.add_mutually_exclusive_group(required=True)
g.add_argument("--input", type=str, help="path to a UTF-8 text file to train on")
g.add_argument("--text", type=str, help="inline training text")
p.add_argument("--output", "-o", type=str, required=True,
help="output directory for the saved model")
p.add_argument("--preset", type=str, default=None,
choices=list(PRESETS.keys()),
help="use a capacity preset (tiny/small/medium/large/1b)")
p.add_argument("--D", type=int, default=None, help="hypervector dimensionality")
p.add_argument("--context-window", type=int, default=None)
p.add_argument("--kernel-radius", type=int, default=None)
p.add_argument("--temperature", type=float, default=0.5)
p.add_argument("--push-to-hub", type=str, default=None,
help="push to HF Hub repo (e.g. user/palimpseste-small)")
p.add_argument("--hub-token", type=str, default=None)
p.add_argument("--seed", type=int, default=42)
return p.parse_args()
def build_config(args: argparse.Namespace) -> PalimpsesteConfig:
if args.preset:
cfg = PRESETS[args.preset]
# allow overrides
if args.D is not None:
cfg = PalimpsesteConfig(D=args.D, context_window=cfg.context_window,
kernel_radius=cfg.kernel_radius)
if args.context_window is not None:
cfg.context_window = args.context_window
if args.kernel_radius is not None:
cfg.kernel_radius = args.kernel_radius
cfg.temperature = args.temperature
return cfg
# fully custom
D = args.D or 10_000
return PalimpsesteConfig(
D=D,
context_window=args.context_window or 32,
kernel_radius=args.kernel_radius or 200,
temperature=args.temperature,
)
def load_text(args: argparse.Namespace) -> str:
if args.text is not None:
return args.text
p = Path(args.input)
if not p.exists():
sys.exit(f"error: input file not found: {p}")
return p.read_text(encoding="utf-8")
def main() -> None:
args = parse_args()
config = build_config(args)
text = load_text(args)
print("=" * 70)
print("PALIMPSESTE — causal LM training")
print("=" * 70)
print(f"config: D={config.D:,} context_window={config.context_window} "
f"kernel_radius={config.kernel_radius} temperature={config.temperature}")
cap_log2 = float(np.log2(0.14 * config.D) + config.D / 4.0)
print(f"theoretical capacity: 2^{cap_log2:.1f} "
f"≈ 10^{cap_log2 * 0.30103:.1f} associations")
print(f"corpus: {len(text):,} chars")
print()
rng = np.random.default_rng(args.seed)
lm = HFPalimpsesteLM(config=config, rng=rng)
# build tokenizer
t0 = time.perf_counter()
lm.build_tokenizer(text)
print(f"vocab built: {lm.tokenizer.vocab_size} chars "
f"({time.perf_counter() - t0:.2f}s)")
# train (O(1) per token)
t0 = time.perf_counter()
n_tokens = lm.train_on_text(text, verbose=True)
dt = time.perf_counter() - t0
print(f"\ntrained: {n_tokens:,} tokens in {dt:.2f}s "
f"({n_tokens / max(dt, 1e-9):,.0f} tok/s)")
print(f"|M| = {len(lm.mem):,} traces")
# quick self-eval
ev = lm.evaluate(text)
print(f"\nself-eval: next_token_accuracy={ev['next_token_accuracy']:.3f} "
f"mean_surprise={ev['mean_surprise']:.3f}")
# sample generation
sample_prompts = ["", "the ", "a "]
print("\n--- sample generation ---")
for prompt in sample_prompts:
out = lm.generate(prompt, max_new_tokens=40, temperature=0.5, seed=args.seed)
print(f" {prompt!r:8s} -> {out.text!r}")
# save
out_dir = Path(args.output)
print(f"\nsaving to {out_dir} ...")
lm.save_pretrained(out_dir)
print(f"saved. files: {sorted(p.name for p in out_dir.iterdir())}")
# push to hub
if args.push_to_hub:
print(f"\npushing to HF Hub: {args.push_to_hub} ...")
url = lm.push_to_hub(args.push_to_hub, token=args.hub_token)
print(f"pushed: {url}")
print("\n" + "=" * 70)
print("done. Load with:")
print(f" from palimseste.hf import HFPalimpsesteLM")
print(f" lm = HFPalimpsesteLM.from_pretrained('{out_dir}')")
print("=" * 70)
if __name__ == "__main__":
main()
|