JULIAN-100M

pip install jax[cpu] flax orbax-checkpoint sentencepiece

import jax
import jax.numpy as jnp
import sentencepiece as spm
import orbax.checkpoint as ocp
import os

# Add model code to path
import sys
sys.path.insert(0, "path/to/julian-100m")

from src.model import JULIAN_100M, create_model

# Load tokenizer
tokenizer = spm.SentencePieceProcessor()
tokenizer.Load("julian-100m/tokenizer/julian_24k.model")

# Create model
model = create_model(JULIAN_100M)

# Initialize with dummy input
rng = jax.random.PRNGKey(0)
dummy = jnp.ones((1, JULIAN_100M.max_seq_len), dtype=jnp.int32)
variables = model.init(rng, dummy)

# Load checkpoint
checkpoint_path = os.path.abspath("julian-100m/checkpoint")
checkpointer = ocp.StandardCheckpointer()
restored = checkpointer.restore(
    checkpoint_path,
    args=ocp.args.StandardRestore({"params": variables["params"]})
)
params = restored["params"]

print("✅ Model loaded successfully!")
print(f"Parameters: {sum(p.size for p in jax.tree_util.tree_leaves(params)):,}")