| |
| |
| |
| import optax |
|
|
| def load_pile_streaming(batch_size, seq_len, seed=42): |
| from datasets import load_dataset |
| ds = load_dataset("monology/pile-uncopyrighted", streaming=True, split="train") |
| def tokenize(example): |
| text = example.get("text", "") |
| if not text: text = "" |
| tokens = list(text.encode('utf-8'))[:seq_len] |
| tokens = tokens + [0] * (seq_len - len(tokens)) |
| return {"input_ids": tokens} |
| ds = ds.map(tokenize) |
| ds = ds.shuffle(seed=seed, buffer_size=1000) |
| return ds |
|
|
| def create_batches(ds, batch_size, seq_len, d_model, n_batches=50): |
| batches = [] |
| it = iter(ds) |
| for _ in range(n_batches): |
| batch = [] |
| for _ in range(batch_size): |
| try: |
| example = next(it) |
| ids = example.get("input_ids", [0]*seq_len) |
| if len(ids) < seq_len: ids = ids + [0]*(seq_len-len(ids)) |
| batch.append(ids[:seq_len]) |
| except StopIteration: batch.append([0]*seq_len) |
| ids_arr = np.array(batch, dtype=np.float32) |
| x = ids_arr / 256.0 |
| x = np.tile(x[..., None], (1,1,d_model//seq_len+1))[:,:,:d_model] |
| x = x * (0.5 + 0.5*np.random.randn(batch_size,seq_len,d_model)*0.1) |
| batches.append(jnp.array(x.astype(np.float32))) |
| return batches |
|
|
| def train_model(model, d_model, n_steps=30, lr=1e-4): |
| """Train CogniThread model on PILE data with wave-scalar computing.""" |
| print(f" Training on PILE (streaming) for {n_steps} steps...") |
| try: |
| ds = load_pile_streaming(batch_size=2, seq_len=8) |
| batches = create_batches(ds, batch_size=2, seq_len=8, d_model=d_model, n_batches=n_steps) |
| print(f" Loaded {len(batches)} batches from PILE") |
| except Exception as e: |
| print(f" PILE loading failed ({e}), using synthetic data") |
| key = jax.random.key(0) |
| batches = [jax.random.normal(jax.random.fold_in(key, i), (2, 8, d_model)) |
| for i in range(n_steps)] |
|
|
| |
| tx = optax.chain( |
| optax.clip_by_global_norm(1.0), |
| optax.adam(lr), |
| ) |
| opt = nn.ModelAndOptimizer(model, tx, wrt=nn.Param) |
|
|
| @nn.jit |
| def train_step(opt, x, step_key): |
| def loss_fn(model): |
| y_final, var_y, w_f, mu = model(x, key=step_key) |
| |
| |
| y_sq = jnp.sum(y_final ** 2) |
| y_sq = jnp.where(jnp.isfinite(y_sq), y_sq, jnp.array(1e6)) |
| var_loss = jnp.mean(var_y) * 0.01 |
| ent_loss = -jnp.mean(mu * jnp.log(mu + 1e-8)) * 0.001 |
| return jnp.mean(y_sq) + var_loss + ent_loss |
| loss, grads = nn.value_and_grad(loss_fn)(opt.model) |
| |
| grads = jax.tree.map(lambda g: jnp.where(jnp.isfinite(g), g, jnp.zeros_like(g)), grads) |
| opt.update(grads) |
| return jnp.where(jnp.isfinite(loss), loss, jnp.array(1e6)) |
|
|
| losses = [] |
| for step, x_batch in enumerate(batches): |
| step_key = jax.random.key(step) |
| loss = train_step(opt, x_batch, step_key) |
| losses.append(float(loss)) |
| if step % 10 == 0 or step == n_steps - 1: |
| print(f" Step {step+1}/{n_steps}: loss={float(loss):.4f}") |
|
|
| print(f" Training complete. Final loss: {losses[-1]:.4f}") |
| return opt.model, losses |
|
|
