dsm.livemd

Diffusion Encoder

System.put_env("XLA_BUILD", "true")

Mix.install([
  {:nx, "~> 0.11.0"},
  {:exla, "~> 0.11.0"},
  {:kino_vega_lite, "~> 0.1.11"}
])

System.put_env("HSA_OVERRIDE_GFX_VERSION", "11.0.0")
System.put_env("XLA_TARGET", "rocm")

# ⚡ 1. Disable the buggy AMD Fusion pass
System.put_env("TF_ROCM_FUSION_ENABLE", "0")

# ⚡ 2. Disable XLA Autotuning completely
System.put_env("XLA_FLAGS", "--xla_gpu_autotune_level=0 --xla_gpu_enable_reassociation_for_converted_ar=false --xla_gpu_graph_enable_concurrent_region=false")

System.put_env("ROCM_DISABLE_CU_MEMSET", "1")

Application.put_env(:exla, :clients,
  rocm: [
    platform: :rocm,
    memory_fraction: 0.75,  # ← force total claimed under 8GB
    preallocate: false,
    num_replicas: 1
  ]
)

System.put_env("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.33")

Application.put_env(:exla, :default_client, :rocm)

gpu = {EXLA.Backend, client: :rocm}
Nx.global_default_backend(gpu)
Nx.Defn.default_options(compiler: EXLA)

platforms = EXLA.Client.get_supported_platforms()
if platforms[:rocm] do
  IO.puts("✅ AMD RX 7600 online!")
else
  IO.puts("❌ GPU missing — check HSA_OVERRIDE_GFX_VERSION")
end

Diffusion

defmodule DLM.Tree do
  @moduledoc """
  Recursively maps a function over nested tuples of tensors.
  Works perfectly both in standard Elixir AND inside Nx `defn`.
  """
  
  def map(tuple, func) when is_tuple(tuple) do
    tuple
    |> Tuple.to_list()
    |> Enum.map(&map(&1, func))
    |> List.to_tuple()
  end

  # The base case: When it hits a tensor (or a compile-time Expr), apply the function
  def map(tensor, func) do
    func.(tensor)
  end
end

{:module, DLM.Tree, <<70, 79, 82, 49, 0, 0, 8, ...>>, {:map, 2}}

defmodule DLM.ShapeDiffusion do
  import Nx.Defn

  # 1. We initialize an embedding table that will be molded by diffusion
  def init_state(vocab_size \\ 128, dim \\ 512, seed \\ 42) do
    key = Nx.Random.key(seed)
    
    # The raw clay: character embeddings
    {embeds, key} = Nx.Random.normal(key, 0.0, 0.1, shape: {vocab_size, dim})
    
    # Denoiser params (simplified MLP for example)
    {d_w1, key} = Nx.Random.normal(key, 0.0, 0.05, shape: {dim, dim * 4})
    {d_w2, key} = Nx.Random.normal(key, 0.0, 0.05, shape: {dim * 4, dim})
    
    {embeds, d_w1, d_w2, key}
  end

  # Standard L2 Norm to keep the space from expanding infinitely
  defn l2_normalize(tensor) do
    norm = Nx.sqrt(Nx.sum(Nx.pow(tensor, 2), axes: [-1], keep_axes: true))
    Nx.divide(tensor, Nx.add(norm, 1.0e-6))
  end

  # The Denoising Network
  defn denoise(x_t, d_w1, d_w2) do
    # Softplus hidden layer
    hidden = Nx.dot(x_t, d_w1)
    hidden = Nx.select(hidden > 20.0, hidden, Nx.log1p(Nx.exp(hidden)))
    
    Nx.dot(hidden, d_w2)
  end
end

{:module, DLM.ShapeDiffusion, <<70, 79, 82, 49, 0, 0, 21, ...>>, true}

defmodule DLM.Phase1Optimizer do
  import Nx.Defn

  def init(params) do
    zero_like = fn t -> Nx.broadcast(Nx.tensor(0.0, type: :f32), t) end
    ms = DLM.Tree.map(params, zero_like)
    vs = DLM.Tree.map(params, zero_like)
    {ms, vs}
  end

  defn newton_schulz(g) do
    # 1. Normalize the gradient matrix by its Frobenius norm
    norm = Nx.sqrt(Nx.sum(Nx.pow(g, 2))) + 1.0e-8
    x = g / norm

    # 2. 5-Step Orthogonal Projection
    {final_x, _} =
      while {mat = x, i = 0}, Nx.less(i, 5) do
        mat_t = Nx.transpose(mat)
        mat_mat_t = Nx.dot(mat, mat_t)
        
        # X = 1.5 * X - 0.5 * (X * X^T) * X
        term1 = mat * 1.5
        term2 = Nx.dot(mat_mat_t, mat) * 0.5
        
        {term1 - term2, i + 1}
      end

    final_x
  end

  defn muon_step(param, grad, m, v, lr) do
    # Muon uses standard momentum rather than Adam's variance-scaled momentum
    beta1 = 0.95
    new_m = beta1 * m + (1.0 - beta1) * grad
    
    # Force the momentum update into a pure rotational matrix
    ortho_grad = newton_schulz(new_m)
    
    # Apply the orthogonal step
    param_update = param - (lr * ortho_grad)
    
    # Return updated param, new momentum, and pass through the untouched v
    {param_update, new_m, v}
  end

  defn adam_step(param, grad, m, v, step, lr) do
    beta1 = 0.9
    beta2 = 0.999
    eps = 1.0e-8
    
    new_m = beta1 * m + (1.0 - beta1) * grad
    new_v = beta2 * v + (1.0 - beta2) * Nx.pow(grad, 2)
    
    m_hat = new_m / (1.0 - Nx.pow(beta1, step + 1.0))
    v_hat = new_v / (1.0 - Nx.pow(beta2, step + 1.0))
    
    param_update = param - lr * m_hat / (Nx.sqrt(v_hat) + eps)
    {param_update, new_m, new_v}
  end
end

{:module, DLM.Phase1Optimizer, <<70, 79, 82, 49, 0, 0, 22, ...>>, true}

defmodule DLM.NoiseSchedule do
  import Nx.Defn

  def build(timesteps \\ 1000) do
    beta_start = 0.0001
    beta_end = 0.02
    betas = Nx.linspace(beta_start, beta_end, n: timesteps)
    alphas = Nx.subtract(1.0, betas)
    
    alphas_culmprod = Nx.tensor(Enum.scan(Nx.to_flat_list(alphas), &(&1 * &2)))
    
    {Nx.backend_copy(betas), Nx.backend_copy(alphas_culmprod)}
  end

  # Move batch_size and timesteps to opts so they are evaluated at compile-time
  defn sample_t(key, alphas_culmprod, opts \\ []) do
    # Define defaults, but we will pass them dynamically from the loop
    opts = keyword!(opts, batch_size: 4096, timesteps: 1000)
    
    batch_size = opts[:batch_size]
    timesteps = opts[:timesteps]

    # Now batch_size is a pure integer, not a tensor!
    {t, next_key} = Nx.Random.randint(key, 0, timesteps - 1, shape: {batch_size})
    alpha_c_t = Nx.take(alphas_culmprod, t) |> Nx.new_axis(-1)
    {t, alpha_c_t, next_key}
  end
end

{:module, DLM.NoiseSchedule, <<70, 79, 82, 49, 0, 0, 18, ...>>, true}

defmodule DLM.Trainer do
  import Nx.Defn

  defn compute_grad_and_step(params, ms, vs, tokens, noise, alpha_culmprod, step, lr) do
    {embeds, d_w1, d_w2, key} = params
    
    {loss, grads} = value_and_grad(params, fn {emb, w1, w2, _k} ->
      # 1. Get the current continuous shape for the characters
      x_0 = Nx.take(emb, Nx.flatten(tokens)) 
            |> Nx.reshape({Nx.axis_size(tokens, 0), 512})
            |> DLM.ShapeDiffusion.l2_normalize() # Keep them on a hypersphere

      # 2. Corrupt with noise (q_sample)
      mean = x_0 * Nx.sqrt(alpha_culmprod)
      variance = noise * Nx.sqrt(1.0 - alpha_culmprod)
      x_t = mean + variance

      # 3. Predict the noise
      pred_noise = DLM.ShapeDiffusion.denoise(x_t, w1, w2)
      
      # 4. The Loss: Mean Squared Error between predicted and actual noise
      Nx.mean(Nx.pow(pred_noise - noise, 2))
    end)

    # ⚡ 1. Unpack grads, ms, and vs correctly (4 elements because 'key' is in params)
    {g_emb, g_w1, g_w2, _gk} = grads
    {m_emb, m_w1, m_w2, m_key} = ms
    {v_emb, v_w1, v_w2, v_key} = vs
    
    # ⚡ 2. Unpack the 3-tuple returned by muon_step
    {up_emb, up_m_emb, up_v_emb} = DLM.Phase1Optimizer.adam_step(embeds, g_emb, m_emb, v_emb, step, lr)
    {up_w1, up_m_w1, up_v_w1} = DLM.Phase1Optimizer.adam_step(d_w1, g_w1, m_w1, v_w1, step, lr)
    {up_w2, up_m_w2, up_v_w2} = DLM.Phase1Optimizer.adam_step(d_w2, g_w2, m_w2, v_w2, step, lr)
    
    # ⚡ 3. Repack params, ms, and vs to pass back to the training loop
    new_params = {up_emb, up_w1, up_w2, key}
    new_ms = {up_m_emb, up_m_w1, up_m_w2, m_key}
    new_vs = {up_v_emb, up_v_w1, up_v_w2, v_key}

    {loss, new_params, new_ms, new_vs}
  end
end

{:module, DLM.Trainer, <<70, 79, 82, 49, 0, 0, 22, ...>>, true}

# Setup Phase 1
batch_size = 4096 
timesteps = 1000
save_interval = 1000 # ⚡ Checkpoint interval

{_betas, alphas_cumprod} = DLM.NoiseSchedule.build(timesteps)
alphas_cumprod = Nx.backend_copy(alphas_cumprod, gpu)

# Initialize
raw_params = DLM.ShapeDiffusion.init_state(128, 512, 42)
raw_opt_state = DLM.Phase1Optimizer.init(raw_params)

# Move to GPU
params = DLM.Tree.map(raw_params, &Nx.backend_copy(&1, gpu))
{ms, vs} = DLM.Tree.map(raw_opt_state, &Nx.backend_copy(&1, gpu))
key = Nx.Random.key(42) |> Nx.backend_copy(gpu)

full_dataset_2d = File.read!("scholar_full_mixed_128.bin") |> :erlang.binary_to_term()

# Flatten your dataset and calculate steps
dataset_1d = Nx.flatten(full_dataset_2d)
total_chars = Nx.axis_size(dataset_1d, 0)
total_steps = div(total_chars, batch_size)

IO.puts("🚀 IGNITION: Phase 1 Joint Shape Training...")

# The Loop (Now capturing the final state)
{final_p, _final_m, _final_v, _final_key} = 
  Enum.reduce(1..total_steps, {params, ms, vs, key}, fn step, {cp, cm, cv, cur_key} ->
    offset = (step - 1) * batch_size
    
    # 1. Grab a batch of isolated characters
    tokens = Nx.slice_along_axis(dataset_1d, offset, batch_size, axis: 0) 
             |> Nx.backend_copy(gpu)

    # 2. Sample random T and pure Gaussian Noise
    {_t, alpha_c_t, key1} = DLM.NoiseSchedule.sample_t(
      cur_key, 
      alphas_cumprod, 
      batch_size: batch_size, 
      timesteps: timesteps
    )
    {noise, next_key} = Nx.Random.normal(key1, 0.0, 1.0, shape: {batch_size, 512})

    # 3. Step the network
    lr = Nx.tensor(1.0e-4, type: :f32) |> Nx.backend_copy(gpu)

    {loss, up_p, up_m, up_v} = 
      DLM.Trainer.compute_grad_and_step(cp, cm, cv, tokens, noise, alpha_c_t, step, lr)

    # 4. Telemetry
    if rem(step, 100) == 0 do
      IO.puts("➡️  Phase 1 Step #{step} | Diffusion Loss (MSE): #{Float.round(Nx.to_number(loss), 4)}")
    end

    # ⚡ 5. Auto-Save Checkpoint
    if rem(step, save_interval) == 0 do
      IO.puts("💾 Auto-Saving Phase 1 Checkpoint...")
      # Bring params back to CPU before serializing
      cpu_p = DLM.Tree.map(up_p, &Nx.backend_copy(&1, Nx.BinaryBackend))
      File.write!("phase1_shape_diffusion_step#{step}.bin", :erlang.term_to_binary(%{params: cpu_p}))
    end

    # Optional but recommended: Force garbage collection to keep RAM stable on long runs
    if rem(step, 100) == 0, do: :erlang.garbage_collect(self())

    {up_p, up_m, up_v, next_key}
  end)

IO.puts("🎉 PHASE 1 TRAINING COMPLETE!")

# ⚡ Save the final completed masterpiece
cpu_final_p = DLM.Tree.map(final_p, &Nx.backend_copy(&1, Nx.BinaryBackend))
File.write!("phase1_shape_diffusion_COMPLETE.bin", :erlang.term_to_binary(%{params: cpu_final_p}))

🚀 IGNITION: Phase 1 Joint Shape Training...
➡️  Phase 1 Step 100 | Diffusion Loss (MSE): 5.6898
➡️  Phase 1 Step 200 | Diffusion Loss (MSE): 5.6314
➡️  Phase 1 Step 300 | Diffusion Loss (MSE): 5.5871
➡️  Phase 1 Step 400 | Diffusion Loss (MSE): 5.4957
➡️  Phase 1 Step 500 | Diffusion Loss (MSE): 5.4055
➡️  Phase 1 Step 600 | Diffusion Loss (MSE): 5.3435
➡️  Phase 1 Step 700 | Diffusion Loss (MSE): 5.2726
➡️  Phase 1 Step 800 | Diffusion Loss (MSE): 5.239
➡️  Phase 1 Step 900 | Diffusion Loss (MSE): 5.1454
➡️  Phase 1 Step 1000 | Diffusion Loss (MSE): 5.0864
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 1100 | Diffusion Loss (MSE): 5.0293
➡️  Phase 1 Step 1200 | Diffusion Loss (MSE): 4.9286
➡️  Phase 1 Step 1300 | Diffusion Loss (MSE): 4.8879
➡️  Phase 1 Step 1400 | Diffusion Loss (MSE): 4.825
➡️  Phase 1 Step 1500 | Diffusion Loss (MSE): 4.7463
➡️  Phase 1 Step 1600 | Diffusion Loss (MSE): 4.7186
➡️  Phase 1 Step 1700 | Diffusion Loss (MSE): 4.6309
➡️  Phase 1 Step 1800 | Diffusion Loss (MSE): 4.5989
➡️  Phase 1 Step 1900 | Diffusion Loss (MSE): 4.5101
➡️  Phase 1 Step 2000 | Diffusion Loss (MSE): 4.4519
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 2100 | Diffusion Loss (MSE): 4.3863
➡️  Phase 1 Step 2200 | Diffusion Loss (MSE): 4.3545
➡️  Phase 1 Step 2300 | Diffusion Loss (MSE): 4.2866
➡️  Phase 1 Step 2400 | Diffusion Loss (MSE): 4.2364
➡️  Phase 1 Step 2500 | Diffusion Loss (MSE): 4.1808
➡️  Phase 1 Step 2600 | Diffusion Loss (MSE): 4.0944
➡️  Phase 1 Step 2700 | Diffusion Loss (MSE): 4.0663
➡️  Phase 1 Step 2800 | Diffusion Loss (MSE): 3.9988
➡️  Phase 1 Step 2900 | Diffusion Loss (MSE): 3.9714
➡️  Phase 1 Step 3000 | Diffusion Loss (MSE): 3.878
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 3100 | Diffusion Loss (MSE): 3.849
➡️  Phase 1 Step 3200 | Diffusion Loss (MSE): 3.786
➡️  Phase 1 Step 3300 | Diffusion Loss (MSE): 3.738
➡️  Phase 1 Step 3400 | Diffusion Loss (MSE): 3.6879
➡️  Phase 1 Step 3500 | Diffusion Loss (MSE): 3.6292
➡️  Phase 1 Step 3600 | Diffusion Loss (MSE): 3.5701
➡️  Phase 1 Step 3700 | Diffusion Loss (MSE): 3.5043
➡️  Phase 1 Step 3800 | Diffusion Loss (MSE): 3.4794
➡️  Phase 1 Step 3900 | Diffusion Loss (MSE): 3.4314
➡️  Phase 1 Step 4000 | Diffusion Loss (MSE): 3.3905
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 4100 | Diffusion Loss (MSE): 3.3543
➡️  Phase 1 Step 4200 | Diffusion Loss (MSE): 3.2889
➡️  Phase 1 Step 4300 | Diffusion Loss (MSE): 3.2413
➡️  Phase 1 Step 4400 | Diffusion Loss (MSE): 3.1919
➡️  Phase 1 Step 4500 | Diffusion Loss (MSE): 3.1653
➡️  Phase 1 Step 4600 | Diffusion Loss (MSE): 3.13
➡️  Phase 1 Step 4700 | Diffusion Loss (MSE): 3.0708
➡️  Phase 1 Step 4800 | Diffusion Loss (MSE): 3.0318
➡️  Phase 1 Step 4900 | Diffusion Loss (MSE): 2.9618
➡️  Phase 1 Step 5000 | Diffusion Loss (MSE): 2.9353
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 5100 | Diffusion Loss (MSE): 2.9003
➡️  Phase 1 Step 5200 | Diffusion Loss (MSE): 2.8522
➡️  Phase 1 Step 5300 | Diffusion Loss (MSE): 2.8091
➡️  Phase 1 Step 5400 | Diffusion Loss (MSE): 2.788
➡️  Phase 1 Step 5500 | Diffusion Loss (MSE): 2.7313
➡️  Phase 1 Step 5600 | Diffusion Loss (MSE): 2.7074
➡️  Phase 1 Step 5700 | Diffusion Loss (MSE): 2.6606
➡️  Phase 1 Step 5800 | Diffusion Loss (MSE): 2.6175
➡️  Phase 1 Step 5900 | Diffusion Loss (MSE): 2.5807
➡️  Phase 1 Step 6000 | Diffusion Loss (MSE): 2.5495
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 6100 | Diffusion Loss (MSE): 2.5131
➡️  Phase 1 Step 6200 | Diffusion Loss (MSE): 2.4793
➡️  Phase 1 Step 6300 | Diffusion Loss (MSE): 2.4427
➡️  Phase 1 Step 6400 | Diffusion Loss (MSE): 2.4165
➡️  Phase 1 Step 6500 | Diffusion Loss (MSE): 2.3734
➡️  Phase 1 Step 6600 | Diffusion Loss (MSE): 2.3547
➡️  Phase 1 Step 6700 | Diffusion Loss (MSE): 2.2991
➡️  Phase 1 Step 6800 | Diffusion Loss (MSE): 2.284
➡️  Phase 1 Step 6900 | Diffusion Loss (MSE): 2.2573
➡️  Phase 1 Step 7000 | Diffusion Loss (MSE): 2.2081
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 7100 | Diffusion Loss (MSE): 2.1595
➡️  Phase 1 Step 7200 | Diffusion Loss (MSE): 2.1416
➡️  Phase 1 Step 7300 | Diffusion Loss (MSE): 2.1135
➡️  Phase 1 Step 7400 | Diffusion Loss (MSE): 2.0869
➡️  Phase 1 Step 7500 | Diffusion Loss (MSE): 2.0512
➡️  Phase 1 Step 7600 | Diffusion Loss (MSE): 2.036
➡️  Phase 1 Step 7700 | Diffusion Loss (MSE): 2.0095
➡️  Phase 1 Step 7800 | Diffusion Loss (MSE): 1.9795
➡️  Phase 1 Step 7900 | Diffusion Loss (MSE): 1.9554
➡️  Phase 1 Step 8000 | Diffusion Loss (MSE): 1.9256
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 8100 | Diffusion Loss (MSE): 1.9041
➡️  Phase 1 Step 8200 | Diffusion Loss (MSE): 1.8744
➡️  Phase 1 Step 8300 | Diffusion Loss (MSE): 1.8548
➡️  Phase 1 Step 8400 | Diffusion Loss (MSE): 1.8408
➡️  Phase 1 Step 8500 | Diffusion Loss (MSE): 1.8061
➡️  Phase 1 Step 8600 | Diffusion Loss (MSE): 1.7823
➡️  Phase 1 Step 8700 | Diffusion Loss (MSE): 1.7664
➡️  Phase 1 Step 8800 | Diffusion Loss (MSE): 1.7419
➡️  Phase 1 Step 8900 | Diffusion Loss (MSE): 1.7115
➡️  Phase 1 Step 9000 | Diffusion Loss (MSE): 1.6993
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 9100 | Diffusion Loss (MSE): 1.6797
➡️  Phase 1 Step 9200 | Diffusion Loss (MSE): 1.6566
➡️  Phase 1 Step 9300 | Diffusion Loss (MSE): 1.6479
➡️  Phase 1 Step 9400 | Diffusion Loss (MSE): 1.6266
➡️  Phase 1 Step 9500 | Diffusion Loss (MSE): 1.6132
➡️  Phase 1 Step 9600 | Diffusion Loss (MSE): 1.5884
➡️  Phase 1 Step 9700 | Diffusion Loss (MSE): 1.5637
➡️  Phase 1 Step 9800 | Diffusion Loss (MSE): 1.5541
➡️  Phase 1 Step 9900 | Diffusion Loss (MSE): 1.5436
➡️  Phase 1 Step 10000 | Diffusion Loss (MSE): 1.5334
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 10100 | Diffusion Loss (MSE): 1.5017
➡️  Phase 1 Step 10200 | Diffusion Loss (MSE): 1.4901
➡️  Phase 1 Step 10300 | Diffusion Loss (MSE): 1.4786
➡️  Phase 1 Step 10400 | Diffusion Loss (MSE): 1.4666
➡️  Phase 1 Step 10500 | Diffusion Loss (MSE): 1.4429
➡️  Phase 1 Step 10600 | Diffusion Loss (MSE): 1.4319
➡️  Phase 1 Step 10700 | Diffusion Loss (MSE): 1.4319
➡️  Phase 1 Step 10800 | Diffusion Loss (MSE): 1.4143
➡️  Phase 1 Step 10900 | Diffusion Loss (MSE): 1.3971
➡️  Phase 1 Step 11000 | Diffusion Loss (MSE): 1.3931
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 11100 | Diffusion Loss (MSE): 1.375
➡️  Phase 1 Step 11200 | Diffusion Loss (MSE): 1.3667
➡️  Phase 1 Step 11300 | Diffusion Loss (MSE): 1.3609
➡️  Phase 1 Step 11400 | Diffusion Loss (MSE): 1.3485
➡️  Phase 1 Step 11500 | Diffusion Loss (MSE): 1.3303
➡️  Phase 1 Step 11600 | Diffusion Loss (MSE): 1.3175
➡️  Phase 1 Step 11700 | Diffusion Loss (MSE): 1.3076
➡️  Phase 1 Step 11800 | Diffusion Loss (MSE): 1.2949
➡️  Phase 1 Step 11900 | Diffusion Loss (MSE): 1.2828
➡️  Phase 1 Step 12000 | Diffusion Loss (MSE): 1.2703
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 12100 | Diffusion Loss (MSE): 1.2601
➡️  Phase 1 Step 12200 | Diffusion Loss (MSE): 1.2411
➡️  Phase 1 Step 12300 | Diffusion Loss (MSE): 1.2336
➡️  Phase 1 Step 12400 | Diffusion Loss (MSE): 1.2209
➡️  Phase 1 Step 12500 | Diffusion Loss (MSE): 1.2137
➡️  Phase 1 Step 12600 | Diffusion Loss (MSE): 1.1996
➡️  Phase 1 Step 12700 | Diffusion Loss (MSE): 1.1823
➡️  Phase 1 Step 12800 | Diffusion Loss (MSE): 1.1752
➡️  Phase 1 Step 12900 | Diffusion Loss (MSE): 1.1643
➡️  Phase 1 Step 13000 | Diffusion Loss (MSE): 1.1493
💾 Auto-Saving Phase 1 Checkpoint...
➡️  Phase 1 Step 13100 | Diffusion Loss (MSE): 1.1429
➡️  Phase 1 Step 13200 | Diffusion Loss (MSE): 1.1245
🎉 PHASE 1 TRAINING COMPLETE!

:ok

# =====================================================================
# Phase 1: Manifold Visualization
# =====================================================================
alias VegaLite, as: Vl

# 1. Extract the learned 128x512 embeddings from your final params tuple
{final_embeds, _w1, _w2, _key} = final_p

# Bring it back to the CPU for LinAlg processing
coords = Nx.backend_copy(final_embeds, Nx.BinaryBackend)

# 2. PCA: Squash 512D down to 2D
mean = Nx.mean(coords, axes: [0], keep_axes: true)
centered = Nx.subtract(coords, mean)

# Perform Singular Value Decomposition
{_u, _s, vt} = Nx.LinAlg.svd(centered)

# Extract the top 2 Principal Components
v_top2 = Nx.transpose(vt)[[.., 0..1]]

# Project the 512D coordinates onto the new 2D plane
coords_2d = Nx.dot(centered, v_top2)

# 3. Extract X and Y for graphing
x_vals = coords_2d[[.., 0]] |> Nx.to_flat_list()
y_vals = coords_2d[[.., 1]] |> Nx.to_flat_list()

# 4. Tag and categorize the printable ASCII characters
plot_data =
  Enum.map(32..126, fn ascii ->
    char = List.to_string([ascii])
    
    type = cond do
      char == " " -> "Space"
      char =~ ~r/[aeiouAEIOU]/ -> "Vowel"
      char =~ ~r/[a-zA-Z]/ -> "Consonant"
      char =~ ~r/[0-9]/ -> "Number"
      true -> "Punctuation"
    end

    %{
      "character" => char,
      "x" => Enum.at(x_vals, ascii),
      "y" => Enum.at(y_vals, ascii),
      "type" => type
    }
  end)

# 5. Render the Latent Map
Vl.new(width: 800, height: 600, title: "Diffusion-Molded 512D Latent Manifold")
|> Vl.data_from_values(plot_data)
|> Vl.mark(:text, size: 16, font: "monospace", font_weight: "bold")
|> Vl.encode_field(:x, "x", type: :quantitative, title: "Principal Component 1")
|> Vl.encode_field(:y, "y", type: :quantitative, title: "Principal Component 2")
|> Vl.encode_field(:text, "character", type: :nominal)
|> Vl.encode_field(:color, "type", type: :nominal, scale: [
  domain: ["Space", "Vowel", "Consonant", "Number", "Punctuation"],
  range: ["#FF0000", "#00AEEF", "#2A363B", "#99B898", "#E84A5F"]
])
|> Kino.VegaLite.new()

21:44:18.976 [info] Merging Dots in computation: region_7.7

21:44:18.976 [info] Merging Dots in computation: region_16.14

21:44:18.976 [info] Merging Dots in computation: region_54.58

21:44:18.976 [info] Merging Dots in computation: region_6.61

Data

defmodule BabyLM.ScholarParser do
  @sequence_length 128 

  @datasets %{
    childes: ["childes.train.txt"], 
    bnc: ["bnc_spoken.train.txt"],
    switchboard: ["switchboard.train.txt"],
    subtitles: ["open_subtitles.train.txt"], 
    wiki: ["simple_wiki.train.txt"],
    gutenberg: ["gutenberg.train.txt"]
  }

  def build_scholar_curriculum(directory_path) do
    childes_chunks = load_and_chunk(directory_path, @datasets.childes)
    bnc_chunks = load_and_chunk(directory_path, @datasets.bnc)
    switchboard_chunks = load_and_chunk(directory_path, @datasets.switchboard)
    subtitles_chunks = load_and_chunk(directory_path, @datasets.subtitles)
    wiki_chunks = load_and_chunk(directory_path, @datasets.wiki)
    gutenberg_chunks = load_and_chunk(directory_path, @datasets.gutenberg)

    IO.puts("\n🧬 Mixing FULL Dataset Timeline...")

    # Smash all the chunks together into one massive list, then shuffle
    master_timeline = 
      (wiki_chunks ++ gutenberg_chunks ++ subtitles_chunks ++ 
       switchboard_chunks ++ childes_chunks ++ bnc_chunks)
      |> Enum.shuffle()

    IO.puts("📦 Packing #{length(master_timeline)} sequences into Nx Tensor...")
    dataset_tensor = 
      Nx.tensor(master_timeline, type: :u8)
      |> Nx.backend_transfer(Nx.BinaryBackend) 
    
    File.write!("scholar_full_mixed_128.bin", :erlang.term_to_binary(dataset_tensor))
    IO.puts("🎉 scholar_full_mixed_128.bin successfully generated!")
  end

  defp normalize_char(c) do
    cond do
      c < 128 -> c              # ASCII passthrough
      
      # Common accented Latin — map to base character
      c in 0xC0..0xC5 -> ?A     # À Á Â Ã Ä Å
      c in 0xC8..0xCB -> ?E     # È É Ê Ë
      c in 0xCC..0xCF -> ?I     # Ì Í Î Ï
      c in 0xD2..0xD6 -> ?O     # Ò Ó Ô Õ Ö
      c in 0xD9..0xDC -> ?U     # Ù Ú Û Ü
      c in 0xE0..0xE5 -> ?a     # à á â ã ä å
      c in 0xE8..0xEB -> ?e     # è é ê ë
      c in 0xEC..0xEF -> ?i     # ì í î ï
      c in 0xF2..0xF6 -> ?o     # ò ó ô õ ö
      c in 0xF9..0xFC -> ?u     # ù ú û ü
      c == 0xD1 -> ?N           # Ñ
      c == 0xF1 -> ?n           # ñ
      c == 0xC7 -> ?C           # Ç
      c == 0xE7 -> ?c           # ç
      c == 0xDF -> ?s           # ß -> ss would be ideal but single char is fine
      c == 0xD8 -> ?O           # Ø
      c == 0xF8 -> ?o           # ø
      c == 0xC6 -> ?A           # Æ
      c == 0xE6 -> ?a           # æ
  
      # Smart quotes -> straight quotes
      c in [0x2018, 0x2019] -> ?'
      c in [0x201C, 0x201D] -> ?"
      
      # Dashes
      c in [0x2013, 0x2014] -> ?-   # en-dash, em-dash
      
      # Ellipsis
      c == 0x2026 -> ?.
      
      # Everything else — use a dedicated replacement token
      # Use ASCII 1 (SOH) which never appears naturally in text
      true -> 1
    end
  end

  defp load_and_chunk(dir, files) do
    raw_text =
      files
      |> Enum.map(fn file ->
        path = Path.join(dir, file)
        if File.exists?(path) do
          IO.puts("  ✅ Loaded #{file}")
          File.read!(path) |> String.replace(~r/\s+/, " ")
        else
          IO.puts("  ❌ Missing #{file} (Skipping...)")
          ""
        end
      end)
      |> Enum.join(" ")
  
    raw_text
    |> String.to_charlist()
    |> Enum.map(&normalize_char/1)
    # Filter out chunks with too many replacement tokens (token 1)
    # rather than too many 63s
    |> Enum.chunk_every(@sequence_length, @sequence_length, :discard)
    |> Enum.filter(fn chunk -> 
      replacement_count = Enum.count(chunk, &(&1 == 1))
      replacement_count < 8   # stricter than before — less than 3% replacements
    end)
  end

  defp sample_chunks(chunks, percentage, _base_amount) do
    take_amount = round(length(chunks) * percentage)
    chunks |> Enum.shuffle() |> Enum.take(take_amount)
  end
end

# === RUN THE PARSER ===
# Point this to the folder containing your new BabyLM text files.
# It will ignore everything else in the folder.
data_dir = "./" 

BabyLM.ScholarParser.build_scholar_curriculum(data_dir)