mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/gemma3_text.rb

@@ -0,0 +1,270 @@
+require_relative "cache"
+require_relative "rope_utils"
+
+module MlxLm
+  module Models
+    module Gemma3Text
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "gemma3_text"
+        field :hidden_size, default: 1152
+        field :num_hidden_layers, default: 26
+        field :intermediate_size, default: 6912
+        field :num_attention_heads, default: 4
+        field :head_dim, default: 256
+        field :rms_norm_eps, default: 1.0e-6
+        field :vocab_size, default: 262144
+        field :num_key_value_heads, default: 1
+        field :rope_theta, default: 1_000_000.0
+        field :rope_local_base_freq, default: 10_000.0
+        field :query_pre_attn_scalar, default: 256.0
+        field :sliding_window, default: 512
+        field :sliding_window_pattern, default: 6
+        field :max_position_embeddings, default: 32768
+        field :rope_scaling, default: nil
+      end
+
+      class RMSNorm < MLX::NN::Module
+        def initialize(dims:, eps: 1e-6)
+          super()
+          self.weight = MLX::Core.ones([dims])
+          @eps = eps
+        end
+
+        def call(x)
+          mx = MLX::Core
+          x_sq = x * x
+          mean_sq = mx.mean(x_sq, -1, keepdims: true)
+          x * mx.rsqrt(mean_sq + @eps) * (1.0 + weight)
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args, layer_idx)
+          super()
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = args.query_pre_attn_scalar**(-0.5)
+          pattern = [args.sliding_window_pattern.to_i, 1].max
+          @is_sliding = ((layer_idx + 1) % pattern) != 0
+
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: false)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: false)
+
+          self.q_norm = RMSNorm.new(dims: @head_dim, eps: args.rms_norm_eps)
+          self.k_norm = RMSNorm.new(dims: @head_dim, eps: args.rms_norm_eps)
+
+          if @is_sliding
+            self.rope = MlxLm::Models.initialize_rope(
+              @head_dim,
+              args.rope_local_base_freq,
+              false
+            )
+          else
+            self.rope = MlxLm::Models.initialize_rope(
+              @head_dim,
+              args.rope_theta,
+              false,
+              args.rope_scaling,
+              max_position_embeddings: args.max_position_embeddings
+            )
+          end
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x).reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          queries = q_norm.call(queries)
+          keys = k_norm.call(keys)
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(dim, hidden_dim)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: false)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+        end
+
+        def call(x)
+          down_proj.call(MLX::NN.gelu_approx(gate_proj.call(x)) * up_proj.call(x))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args, layer_idx)
+          super()
+          self.self_attn = Attention.new(args, layer_idx)
+          self.mlp = MLP.new(args.hidden_size, args.intermediate_size)
+          self.input_layernorm = RMSNorm.new(dims: args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = RMSNorm.new(dims: args.hidden_size, eps: args.rms_norm_eps)
+          self.pre_feedforward_layernorm = RMSNorm.new(dims: args.hidden_size, eps: args.rms_norm_eps)
+          self.post_feedforward_layernorm = RMSNorm.new(dims: args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = self_attn.call(input_layernorm.call(x), mask: mask, cache: cache)
+          h = clip_residual(x, post_attention_layernorm.call(r))
+          r = mlp.call(pre_feedforward_layernorm.call(h))
+          clip_residual(h, post_feedforward_layernorm.call(r))
+        end
+
+        private
+
+        def clip_residual(x, y)
+          mx = MLX::Core
+          return x + y unless x.dtype == mx.float16
+
+          bound = mx.finfo(mx.float16).max
+          mx.clip(
+            x.astype(mx.float32) + y.astype(mx.float32),
+            -bound,
+            bound
+          ).astype(mx.float16)
+        end
+      end
+
+      class Gemma3Model < MLX::NN::Module
+        attr_reader :sliding_window_pattern
+
+        def initialize(args)
+          super()
+          @args = args
+          @window_size = args.sliding_window
+          @sliding_window_pattern = [args.sliding_window_pattern.to_i, 1].max
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) do |layer_idx|
+            TransformerBlock.new(args, layer_idx)
+          end
+          self.norm = RMSNorm.new(dims: args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          h = input_embeddings || embed_tokens.call(inputs)
+          h = h * Math.sqrt(@args.hidden_size)
+          layer_cache = cache || [nil] * layers.length
+
+          global_idx = sliding_window_pattern - 1
+          global_mask = _create_attention_mask(h, layer_cache[global_idx])
+          sliding_window_mask = if sliding_window_pattern > 1
+            _create_attention_mask(h, layer_cache[0], window_size: @window_size)
+          else
+            nil
+          end
+
+          layers.each_with_index do |layer, i|
+            is_global = (i % sliding_window_pattern) == (sliding_window_pattern - 1)
+            mask = is_global ? global_mask : sliding_window_mask
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache = nil, window_size: nil)
+          n = h.shape[1]
+          return cache.make_mask(n) if cache && cache.respond_to?(:make_mask)
+
+          if window_size
+            offset = cache ? cache.offset : 0
+            if cache && cache.instance_variable_defined?(:@max_size)
+              max_size = cache.instance_variable_get(:@max_size)
+              offset = [max_size - 1, offset].min if max_size && max_size > 0
+            end
+            return _create_causal_mask(n, offset: offset, window_size: window_size) if offset + n > window_size
+          end
+
+          return nil if n == 1
+
+          "causal"
+        end
+
+        def _create_causal_mask(n, offset: 0, window_size: nil)
+          mx = MLX::Core
+          rinds = mx.arange(0, offset + n, 1, mx.int32).reshape([1, offset + n])
+          linds = mx.arange(offset, offset + n, 1, mx.int32).reshape([n, 1])
+
+          mask = mx.greater_equal(linds, rinds)
+          if window_size
+            mask = mx.logical_and(mask, mx.less(linds, mx.add(rinds, window_size)))
+          end
+          mask
+        end
+      end
+
+      class Model < MLX::NN::Module
+        attr_reader :args
+
+        def initialize(args)
+          super()
+          @args = args
+          @tie_word_embeddings = false
+          self.model_type = args.model_type
+          self.model = Gemma3Model.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          out = model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+          if @tie_word_embeddings || lm_head.nil?
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          sanitized = weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+          unless sanitized.key?("lm_head.weight")
+            @tie_word_embeddings = true
+            self.lm_head = nil
+          end
+          sanitized
+        end
+
+        def layers
+          model.layers
+        end
+
+        def make_cache
+          pattern = [@args.sliding_window_pattern.to_i, 1].max
+          max_size = @args.sliding_window || @args.max_position_embeddings || 1
+          Array.new(@args.num_hidden_layers) do |i|
+            is_global = (i % pattern) == (pattern - 1)
+            if is_global
+              MlxLm::KVCache.new
+            else
+              MlxLm::RotatingKVCache.new(max_size: max_size, keep: 0)
+            end
+          end
+        end
+      end
+
+      Models.register("gemma3_text", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/gemma3n.rb

@@ -0,0 +1,79 @@
+require_relative "gemma2"
+
+module MlxLm
+  module Models
+    module Gemma3n
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "gemma3n"
+        field :text_config, default: nil
+
+        def self.from_dict(params)
+          has_text_config = params.key?("text_config") || params.key?(:text_config)
+          return super if has_text_config
+
+          new(model_type: params["model_type"] || params[:model_type], text_config: params)
+        end
+
+        def initialize(**kwargs)
+          super
+          @text_config = (@text_config || {}).dup
+        end
+      end
+
+      class Model < MLX::NN::Module
+        MULTIMODAL_MODEL_PREFIXES = %w[
+          model.vision_tower
+          model.audio_tower
+          model.embed_audio
+          model.embed_vision
+        ].freeze
+
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.language_model = Gemma2::Model.new(Gemma2::ModelArgs.from_dict(_text_config_for_gemma2(args)))
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          supports_input_embeddings = language_model.method(:call).parameters.any? do |_, name|
+            name == :input_embeddings
+          end
+
+          if supports_input_embeddings
+            language_model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+          else
+            language_model.call(inputs, cache: cache)
+          end
+        end
+
+        def sanitize(weights)
+          weights.reject do |key, _|
+            MULTIMODAL_MODEL_PREFIXES.any? { |prefix| key == prefix || key.start_with?("#{prefix}.") }
+          end
+        end
+
+        def layers
+          language_model.layers
+        end
+
+        def make_cache
+          return language_model.make_cache if language_model.respond_to?(:make_cache)
+
+          nil
+        end
+
+        private
+
+        def _text_config_for_gemma2(args)
+          config = {}
+          (args.text_config || {}).each { |key, value| config[key.to_s] = value }
+          config["model_type"] ||= args.model_type
+          config
+        end
+      end
+
+      Models.register("gemma3n", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/glm.rb

@@ -0,0 +1,164 @@
+module MlxLm
+  module Models
+    module GLM
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "glm"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 28
+        field :intermediate_size, default: 13696
+        field :num_attention_heads, default: 32
+        field :rms_norm_eps, default: 1e-5
+        field :vocab_size, default: 151552
+        field :head_dim, default: nil
+        field :num_key_value_heads, default: nil
+        field :max_position_embeddings, default: nil
+        field :attention_bias, default: false
+        field :rope_theta, default: 10_000.0
+        field :tie_word_embeddings, default: true
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = @head_dim**(-0.5)
+
+          bias = args.attention_bias
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: bias)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: false)
+          self.rope = MLX::NN::RoPE.new(@head_dim, traditional: true, base: args.rope_theta)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x).reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.gate_up_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            2 * args.intermediate_size,
+            bias: false
+          )
+          self.down_proj = MLX::NN::Linear.new(
+            args.intermediate_size,
+            args.hidden_size,
+            bias: false
+          )
+        end
+
+        def call(x)
+          mx = MLX::Core
+          x = gate_up_proj.call(x)
+          split_dim = x.shape[-1] / 2
+          gate, up = mx.split(x, [split_dim], -1)
+          down_proj.call(Activations.swiglu(gate, up))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args)
+          self.input_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = self_attn.call(input_layernorm.call(x), mask: mask, cache: cache)
+          h = x + r
+          r = mlp.call(post_attention_layernorm.call(h))
+          h + r
+        end
+      end
+
+      class GLMModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { TransformerBlock.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          @model_type = args.model_type
+          self.model = GLMModel.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          result = weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("glm", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/glm4.rb

@@ -0,0 +1,180 @@
+module MlxLm
+  module Models
+    module GLM4
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "glm4"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 40
+        field :intermediate_size, default: 13696
+        field :num_attention_heads, default: 32
+        field :attention_bias, default: false
+        field :head_dim, default: nil
+        field :rms_norm_eps, default: 1e-5
+        field :vocab_size, default: 151552
+        field :num_key_value_heads, default: nil
+        field :partial_rotary_factor, default: 0.5
+        field :rope_theta, default: 10_000.0
+        field :rope_traditional, default: true
+        field :max_position_embeddings, default: 32768
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+        end
+      end
+
+      class GLM4MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.gate_up_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            2 * args.intermediate_size,
+            bias: false
+          )
+          self.down_proj = MLX::NN::Linear.new(
+            args.intermediate_size,
+            args.hidden_size,
+            bias: false
+          )
+        end
+
+        def call(x)
+          mx = MLX::Core
+          x = gate_up_proj.call(x)
+          split_dim = x.shape[-1] / 2
+          gate, up_states = mx.split(x, [split_dim], -1)
+          down_proj.call(Activations.swiglu(gate, up_states))
+        end
+      end
+
+      class GLM4Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          @head_dim = args.head_dim
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(
+            dim,
+            args.num_attention_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.k_proj = MLX::NN::Linear.new(
+            dim,
+            args.num_key_value_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.v_proj = MLX::NN::Linear.new(
+            dim,
+            args.num_key_value_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.o_proj = MLX::NN::Linear.new(
+            args.num_attention_heads * @head_dim,
+            args.hidden_size,
+            bias: false
+          )
+
+          self.rope = MLX::NN::RoPE.new(
+            (args.partial_rotary_factor * @head_dim).to_i,
+            base: args.rope_theta,
+            traditional: args.rope_traditional
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x).reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class GLM4DecoderLayer < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = GLM4Attention.new(args)
+          self.mlp = GLM4MLP.new(args)
+          self.input_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_self_attn_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_mlp_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          x = x + post_self_attn_layernorm.call(
+            self_attn.call(input_layernorm.call(x), mask: mask, cache: cache)
+          )
+          residual = x
+          post_mlp_layernorm.call(mlp.call(post_attention_layernorm.call(x))) + residual
+        end
+      end
+
+      class GLM4Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { GLM4DecoderLayer.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.model = GLM4Model.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          lm_head.call(out)
+        end
+
+        def sanitize(weights)
+          weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("glm4", Model, ModelArgs)
+    end
+  end
+end