File size: 4,996 Bytes
4bf2f60 | 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 | #!/usr/bin/env python
"""Train the 13M ContinuousThoughtEngine on quality data.
117 tok/s on CPU. 500k quality tokens × 3 epochs = ~3.5 hours.
After training, the engine produces coherent text + accurate embeddings
for the RAG system (retrieval + online learning).
"""
import argparse, math, os, sys, time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch, torch.nn as nn, torch.nn.functional as F
from fractus.continuous_engine import ContinuousThoughtEngine
from fractus.tokenizer import FractusTokenizer
HF_TOKEN = os.environ.get("HF_TOKEN", "")
HF_REPO = "thefinalboss/Fractus-1B"
def save_ckpt(engine, step, loss, acc, d):
os.makedirs(d, exist_ok=True)
p = os.path.join(d, f"fractus_13m_step_{step}.pt")
torch.save({"model_state": engine.state_dict(),
"config": {"d_model": engine.d_model, "vocab_size": engine.vocab_size},
"step": step, "loss": loss, "accuracy": acc}, p)
print(f" [ckpt] {p} ({os.path.getsize(p)/1e6:.0f}MB)", flush=True)
if HF_TOKEN:
try:
from huggingface_hub import HfApi
api = HfApi(token=HF_TOKEN)
api.upload_file(path_or_fileobj=p,
path_in_repo=f"checkpoints/fractus_13m_step_{step}.pt",
repo_id=HF_REPO, repo_type="model")
print(f" [HF] Uploaded", flush=True)
except Exception as e:
print(f" [HF] Skipped: {type(e).__name__}", flush=True)
def main():
p = argparse.ArgumentParser()
p.add_argument("--epochs", type=int, default=3)
p.add_argument("--chunk-len", type=int, default=16)
p.add_argument("--lr", type=float, default=3e-4)
args = p.parse_args()
torch.set_num_threads(os.cpu_count() or 6)
torch.manual_seed(42)
# Load quality data.
corpus = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"data", "quality_500k.pt")
tokens = torch.load(corpus, weights_only=False).long()
print(f"Corpus: {len(tokens):,} quality tokens", flush=True)
# Build engine.
engine = ContinuousThoughtEngine(
vocab_size=50257, d_model=128, n_heads=2, d_head=64,
n_levels=2, n_oscillators=8, coupling_rank=4,
n_experts=4, top_k=2, expert_d_ff=128, siren_rank=32,
)
print(f"Engine: {sum(p.numel() for p in engine.parameters())/1e6:.1f}M params", flush=True)
# Trainer.
from fractus.train.online import OnlineTrainer
trainer = OnlineTrainer(engine, lr=args.lr)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
trainer.optimizer, T_max=args.epochs, eta_min=1e-5)
ckpt_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"checkpoints")
print(f"\nTraining {args.epochs} epochs, chunk={args.chunk_len}", flush=True)
print("=" * 70, flush=True)
initial_loss = None
for epoch in range(args.epochs):
t0 = time.perf_counter()
m = trainer.train_on_stream_chunked(tokens, chunk_len=args.chunk_len)
elapsed = time.perf_counter() - t0
tps = len(tokens) / elapsed
if initial_loss is None:
initial_loss = m["avg_loss"]
scheduler.step()
ppl = math.exp(min(m["avg_loss"], 20))
print(f"Epoch {epoch+1}/{args.epochs}: loss={m['avg_loss']:.3f} ppl={ppl:.1f} "
f"acc={m['accuracy']:.1%} {tps:.0f} tok/s {elapsed/60:.1f}min", flush=True)
# Checkpoint + sample.
save_ckpt(engine, (epoch+1) * (len(tokens) // args.chunk_len),
m["avg_loss"], m["accuracy"], ckpt_dir)
# Generate sample every epoch.
engine.eval()
engine.reset_thought(1)
tok = FractusTokenizer.gpt2_compatible()
prompt = "def fibonacci"
for tid in tok.encode(prompt)[:16]:
engine.tick(torch.tensor([tid]))
generated = list(tok.encode(prompt)[:16])
for _ in range(30):
logits, _ = engine.tick()
generated.append(logits.argmax(dim=-1).item())
print(f" Sample: {tok.decode(generated)[:100]}", flush=True)
# Final coherence test.
print("\n=== COHERENCE TEST ===", flush=True)
tok = FractusTokenizer.gpt2_compatible()
engine.eval()
for prompt in ["What is Python?", "Explain machine learning", "def sort", "The sun is"]:
engine.reset_thought(1)
ids = tok.encode(prompt)[:16]
for tid in ids:
engine.tick(torch.tensor([tid]))
gen = list(ids)
for _ in range(40):
logits, _ = engine.tick()
l = logits[0] / 0.8
tv, ti = l.topk(40)
probs = F.softmax(tv, dim=-1)
gen.append(ti[torch.multinomial(probs, 1).item()].item())
print(f" [{prompt}] -> {tok.decode(gen)[:120]}", flush=True)
print(f"\nDone. Loss: {initial_loss:.3f} -> {m['avg_loss']:.3f}", flush=True)
if __name__ == "__main__":
main()
|