Upload src/model/julian.py with huggingface_hub

src/model/julian.py

@@ -0,0 +1,225 @@
+"""
+Julian Model - 250M Parameter LLM.
+GPT-style decoder-only transformer with modern improvements.
+"""
+
+from typing import Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+from flax.linen import initializers
+
+from .config import JulianConfig
+from .layers import RMSNorm, TransformerBlock, precompute_rope_frequencies
+
+
+class JulianModel(nn.Module):
+    """
+    Julian Language Model.
+
+    A GPT-style decoder-only transformer with:
+    - RMSNorm (instead of LayerNorm)
+    - RoPE positional encoding
+    - SwiGLU activation
+    - No bias terms
+    """
+
+    config: JulianConfig
+
+    def setup(self):
+        config = self.config
+
+        # Token embeddings
+        self.embed_tokens = nn.Embed(
+            num_embeddings=config.vocab_size,
+            features=config.d_model,
+            embedding_init=initializers.normal(config.initializer_range),
+            name="embed_tokens",
+        )
+
+        # Transformer blocks
+        self.layers = [
+            TransformerBlock(config, name=f"layers_{i}")
+            for i in range(config.n_layers)
+        ]
+
+        # Final norm
+        self.norm = RMSNorm(config.d_model, config.rms_norm_eps, name="norm")
+
+        # Output projection (tied with embeddings)
+        self.lm_head = nn.Dense(
+            config.vocab_size,
+            use_bias=False,
+            kernel_init=initializers.normal(config.initializer_range),
+            name="lm_head",
+        )
+
+    def __call__(
+        self,
+        input_ids: jnp.ndarray,
+        deterministic: bool = True,
+    ) -> jnp.ndarray:
+        """
+        Forward pass.
+
+        Args:
+            input_ids: Token IDs [batch, seq_len]
+            deterministic: If True, disable dropout
+
+        Returns:
+            logits: [batch, seq_len, vocab_size]
+        """
+        config = self.config
+        batch_size, seq_len = input_ids.shape
+
+        # Token embeddings
+        hidden_states = self.embed_tokens(input_ids)
+
+        # Precompute RoPE frequencies
+        sin, cos = precompute_rope_frequencies(
+            config.head_dim,
+            config.max_seq_len,
+            config.rope_theta,
+        )
+
+        # Causal mask
+        mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=bool))
+        mask = mask[None, None, :, :]  # [1, 1, seq, seq]
+
+        # Apply transformer layers
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                sin,
+                cos,
+                mask,
+                deterministic,
+            )
+
+        # Final norm
+        hidden_states = self.norm(hidden_states)
+
+        # Project to vocabulary
+        logits = self.lm_head(hidden_states)
+
+        return logits
+
+    def generate(
+        self,
+        input_ids: jnp.ndarray,
+        max_new_tokens: int = 100,
+        temperature: float = 1.0,
+        top_k: Optional[int] = 50,
+        top_p: Optional[float] = 0.9,
+        rng: Optional[jax.random.PRNGKey] = None,
+    ) -> jnp.ndarray:
+        """
+        Generate text autoregressively.
+
+        Args:
+            input_ids: Prompt token IDs [batch, seq_len]
+            max_new_tokens: Maximum tokens to generate
+            temperature: Sampling temperature
+            top_k: Top-k filtering
+            top_p: Nucleus sampling threshold
+            rng: Random key for sampling
+
+        Returns:
+            Generated token IDs [batch, seq_len + max_new_tokens]
+        """
+        if rng is None:
+            rng = jax.random.PRNGKey(0)
+
+        config = self.config
+
+        for _ in range(max_new_tokens):
+            # Truncate if exceeding max length
+            if input_ids.shape[1] >= config.max_seq_len:
+                context = input_ids[:, -config.max_seq_len:]
+            else:
+                context = input_ids
+
+            # Forward pass
+            logits = self(context, deterministic=True)
+
+            # Get next token logits
+            next_logits = logits[:, -1, :] / temperature
+
+            # Top-k filtering
+            if top_k is not None:
+                top_k_logits, top_k_indices = jax.lax.top_k(next_logits, top_k)
+                next_logits = jnp.full_like(next_logits, -1e9)
+                next_logits = next_logits.at[
+                    jnp.arange(next_logits.shape[0])[:, None],
+                    top_k_indices
+                ].set(top_k_logits)
+
+            # Top-p (nucleus) filtering
+            if top_p is not None:
+                sorted_indices = jnp.argsort(next_logits, axis=-1)[:, ::-1]
+                sorted_logits = jnp.take_along_axis(next_logits, sorted_indices, axis=-1)
+                cumprobs = jnp.cumsum(jax.nn.softmax(sorted_logits, axis=-1), axis=-1)
+
+                # Remove tokens with cumulative prob > top_p
+                sorted_mask = cumprobs > top_p
+                sorted_mask = jnp.concatenate([
+                    jnp.zeros_like(sorted_mask[:, :1]),
+                    sorted_mask[:, :-1]
+                ], axis=-1)
+
+                sorted_logits = jnp.where(sorted_mask, -1e9, sorted_logits)
+                next_logits = jnp.take_along_axis(
+                    sorted_logits,
+                    jnp.argsort(sorted_indices, axis=-1),
+                    axis=-1
+                )
+
+            # Sample
+            rng, sample_rng = jax.random.split(rng)
+            probs = jax.nn.softmax(next_logits, axis=-1)
+            next_token = jax.random.categorical(sample_rng, jnp.log(probs + 1e-10))
+            next_token = next_token[:, None]
+
+            # Append
+            input_ids = jnp.concatenate([input_ids, next_token], axis=1)
+
+            # Stop at EOS
+            if jnp.all(next_token == config.eos_token_id):
+                break
+
+        return input_ids
+
+
+def create_model(config: Optional[JulianConfig] = None) -> JulianModel:
+    """Create Julian model instance."""
+    if config is None:
+        config = JulianConfig()
+    return JulianModel(config)
+
+
+def count_params(params) -> int:
+    """Count total parameters in pytree."""
+    return sum(x.size for x in jax.tree_util.tree_leaves(params))
+
+
+if __name__ == "__main__":
+    # Test model creation
+    config = JulianConfig()
+    model = create_model(config)
+
+    # Initialize with dummy input
+    rng = jax.random.PRNGKey(0)
+    dummy_input = jnp.ones((1, 128), dtype=jnp.int32)
+
+    variables = model.init(rng, dummy_input)
+    params = variables["params"]
+
+    n_params = count_params(params)
+    print(f"Julian Model initialized!")
+    print(f"  Config estimate: {config.estimate_params():,}")
+    print(f"  Actual params: {n_params:,} ({n_params/1e6:.1f}M)")
+
+    # Test forward pass
+    logits = model.apply(variables, dummy_input)
+    print(f"  Output shape: {logits.shape}")  # [1, 128, 24000]