mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/rope_utils.rb

@@ -0,0 +1,316 @@
+module MlxLm
+  module Models
+    class SuScaledRoPE < MLX::NN::Module
+      def initialize(
+        dims,
+        base: 10_000.0,
+        max_position_embeddings: 131_072,
+        original_max_position_embeddings: 4096,
+        short_factor: 1.0,
+        long_factor: 1.0,
+        short_mscale: nil,
+        long_mscale: nil
+      )
+        super()
+        mx = MLX::Core
+        @dim = dims
+        @original_max_position_embeddings = original_max_position_embeddings
+
+        freqs = mx.power(
+          base.to_f,
+          mx.divide(mx.arange(0, dims, 2, mx.float32), dims.to_f)
+        )
+        self._freqs = mx.multiply(mx.array(long_factor, dtype: mx.float32), freqs)
+
+        factor = max_position_embeddings.to_f / original_max_position_embeddings
+        self._scale = long_mscale || if factor <= 1.0
+          1.0
+        else
+          Math.sqrt(1 + Math.log(factor) / Math.log(original_max_position_embeddings))
+        end
+      end
+
+      def call(x, offset: 0)
+        mx = MLX::Core
+        x = scale_rotary_part(x, _scale)
+        mx.rope(x, @dim, false, nil, 1.0, offset, _freqs)
+      end
+
+      private
+
+      def scale_rotary_part(x, scale)
+        return x if scale == 1.0
+
+        mx = MLX::Core
+        rotary, rest = mx.split(x, [@dim], -1)
+        mx.concatenate([mx.multiply(rotary, scale), rest], -1)
+      end
+    end
+
+    class Llama3RoPE < MLX::NN::Module
+      def initialize(
+        dims:,
+        max_position_embeddings: 2048,
+        traditional: false,
+        base: 10_000,
+        scaling_config: nil
+      )
+        super()
+        mx = MLX::Core
+
+        @dims = dims
+        @max_position_embeddings = max_position_embeddings
+        @traditional = traditional
+
+        factor = config_value(scaling_config, "factor")
+        low_freq_factor = config_value(scaling_config, "low_freq_factor", 1.0)
+        high_freq_factor = config_value(scaling_config, "high_freq_factor", 4.0)
+        old_context_len = config_value(
+          scaling_config,
+          "original_max_position_embeddings",
+          8192
+        )
+
+        low_freq_wavelen = old_context_len.to_f / low_freq_factor
+        high_freq_wavelen = old_context_len.to_f / high_freq_factor
+
+        freqs = mx.power(
+          base.to_f,
+          mx.divide(mx.arange(0, dims, 2), dims.to_f)
+        )
+        wavelens = mx.multiply(2.0 * Math::PI, freqs)
+
+        freqs = mx.where(
+          mx.greater(wavelens, low_freq_wavelen),
+          mx.multiply(freqs, factor),
+          freqs
+        )
+
+        is_medium_freq = mx.logical_and(
+          mx.greater(wavelens, high_freq_wavelen),
+          mx.less(wavelens, low_freq_wavelen)
+        )
+
+        smooth_factors = mx.divide(
+          mx.subtract(mx.divide(old_context_len.to_f, wavelens), low_freq_factor),
+          (high_freq_factor - low_freq_factor).to_f
+        )
+
+        smooth_freqs = mx.divide(
+          freqs,
+          mx.add(
+            mx.divide(mx.subtract(1.0, smooth_factors), factor.to_f),
+            smooth_factors
+          )
+        )
+
+        self._freqs = mx.where(is_medium_freq, smooth_freqs, freqs)
+      end
+
+      def extra_repr
+        "#{@dims}, traditional=#{@traditional}, max_position_embeddings=#{@max_position_embeddings}"
+      end
+
+      def call(x, offset: 0)
+        MLX::Core.rope(x, @dims, @traditional, nil, 1.0, offset, _freqs)
+      end
+
+      private
+
+      def config_value(config, key, default = nil)
+        return default if config.nil?
+        return config[key] if config.key?(key)
+
+        config.fetch(key.to_sym, default)
+      end
+    end
+
+    class YarnRoPE < MLX::NN::Module
+      def initialize(
+        dims,
+        traditional: false,
+        max_position_embeddings: 2048,
+        base: 10_000,
+        scaling_factor: 1.0,
+        original_max_position_embeddings: 4096,
+        beta_fast: 32,
+        beta_slow: 1,
+        mscale: 1,
+        mscale_all_dim: 0
+      )
+        super()
+        mx = MLX::Core
+
+        self.mscale = yarn_get_mscale(scaling_factor, mscale) /
+                      yarn_get_mscale(scaling_factor, mscale_all_dim)
+
+        freq_extra = mx.power(
+          base.to_f,
+          mx.divide(mx.arange(0, dims, 2, mx.float32), dims.to_f)
+        )
+        freq_inter = mx.multiply(scaling_factor.to_f, freq_extra)
+
+        low, high = yarn_find_correction_range(
+          dims,
+          base,
+          original_max_position_embeddings,
+          beta_fast,
+          beta_slow
+        )
+
+        freq_mask = mx.subtract(1.0, yarn_linear_ramp_mask(low, high, dims / 2))
+        self._freqs = mx.divide(
+          mx.multiply(freq_inter, freq_extra),
+          mx.add(
+            mx.multiply(freq_inter, freq_mask),
+            mx.multiply(freq_extra, mx.subtract(1.0, freq_mask))
+          )
+        )
+
+        @dims = dims
+        @traditional = traditional
+      end
+
+      def call(x, offset: 0)
+        mx = MLX::Core
+        x = scale_rotary_part(x, mscale) unless mscale == 1.0
+
+        mx.rope(x, @dims, @traditional, nil, 1.0, offset, _freqs)
+      end
+
+      private
+
+      def scale_rotary_part(x, scale)
+        mx = MLX::Core
+        rotary, rest = mx.split(x, [@dims], -1)
+        mx.concatenate([mx.multiply(rotary, scale), rest], -1)
+      end
+
+      def yarn_find_correction_dim(dims, base, original_max_position_embeddings, num_rotations)
+        dims * Math.log(original_max_position_embeddings.to_f / (num_rotations * 2 * Math::PI)) /
+          (2 * Math.log(base))
+      end
+
+      def yarn_find_correction_range(dims, base, original_max_position_embeddings, beta_fast, beta_slow)
+        low = yarn_find_correction_dim(dims, base, original_max_position_embeddings, beta_fast).floor
+        high = yarn_find_correction_dim(dims, base, original_max_position_embeddings, beta_slow).ceil
+        [
+          [low, 0].max,
+          [high, dims - 1].min,
+        ]
+      end
+
+      def yarn_get_mscale(scale = 1, mscale = 1)
+        return 1.0 if scale <= 1
+
+        0.1 * mscale * Math.log(scale) + 1.0
+      end
+
+      def yarn_linear_ramp_mask(min_val, max_val, dim)
+        mx = MLX::Core
+
+        max_val += 0.001 if min_val == max_val
+
+        linear = mx.divide(
+          mx.subtract(mx.arange(0, dim, 1, mx.float32), min_val),
+          max_val - min_val
+        )
+        mx.clip(linear, 0.0, 1.0)
+      end
+    end
+
+    module_function
+
+    def initialize_rope(
+      dims,
+      base,
+      traditional,
+      scaling_config = nil,
+      max_position_embeddings: nil
+    )
+      rope_type = if scaling_config
+        config_value(scaling_config, "type") ||
+          config_value(scaling_config, "rope_type", "default")
+      else
+        "default"
+      end
+
+      case rope_type
+      when "default", "linear"
+        scale = rope_type == "linear" ? 1.0 / config_value(scaling_config, "factor") : 1.0
+        MLX::NN::RoPE.new(dims, traditional: traditional, base: base, scale: scale)
+      when "llama3"
+        Llama3RoPE.new(
+          dims: dims,
+          max_position_embeddings: max_position_embeddings,
+          traditional: traditional,
+          base: base,
+          scaling_config: scaling_config
+        )
+      when "yarn", "deepseek_yarn", "telechat3-yarn"
+        rope_kwargs = {}
+        %w[
+          original_max_position_embeddings
+          beta_fast
+          beta_slow
+          mscale
+          mscale_all_dim
+        ].each do |key|
+          value = config_value(scaling_config, key)
+          rope_kwargs[key.to_sym] = value unless value.nil?
+        end
+
+        YarnRoPE.new(
+          dims,
+          max_position_embeddings: max_position_embeddings,
+          traditional: traditional,
+          scaling_factor: config_value(scaling_config, "factor"),
+          base: base,
+          **rope_kwargs
+        )
+      when "longrope"
+        SuScaledRoPE.new(
+          dims,
+          base: base,
+          max_position_embeddings: max_position_embeddings,
+          original_max_position_embeddings: config_value(
+            scaling_config,
+            "original_max_position_embeddings"
+          ),
+          short_factor: config_value(scaling_config, "short_factor"),
+          long_factor: config_value(scaling_config, "long_factor")
+        )
+      when "mrope"
+        mrope_section = config_value(scaling_config, "mrope_section", [])
+        unless mrope_section.length == 3
+          raise ArgumentError,
+            "MRoPE currently only supports 3 sections, got #{mrope_section.length}."
+        end
+
+        MLX::NN::RoPE.new(dims, traditional: traditional, base: base)
+      else
+        raise ArgumentError, "Unsupported RoPE type #{rope_type}"
+      end
+    end
+
+    def config_value(config, key, default = nil)
+      return default if config.nil?
+      return config[key] if config.key?(key)
+
+      config.fetch(key.to_sym, default)
+    end
+    private_class_method :config_value
+
+    module RoPEUtils
+      SuScaledRoPE = MlxLm::Models::SuScaledRoPE
+      Llama3RoPE = MlxLm::Models::Llama3RoPE
+      YarnRoPE = MlxLm::Models::YarnRoPE
+
+      module_function
+
+      def initialize_rope(*args, **kwargs)
+        MlxLm::Models.initialize_rope(*args, **kwargs)
+      end
+    end
+  end
+end

data/lib/mlx_lm/models/rwkv7.rb

@@ -0,0 +1,101 @@
+require_relative "recurrent_gemma"
+
+module MlxLm
+  module Models
+    module Rwkv7
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "rwkv7"
+        field :vocab_size
+        field :hidden_size
+        field :intermediate_size
+        field :norm_eps, default: 1e-5
+        field :head_dim
+        field :num_hidden_layers
+        field :a_low_rank_dim, default: nil
+        field :v_low_rank_dim, default: nil
+        field :gate_low_rank_dim, default: nil
+        field :decay_low_rank_dim, default: nil
+        field :tie_word_embeddings, default: false
+        field :rope_theta, default: 10_000.0
+        field :attention_window_size, default: 128
+        field :block_types, default: nil
+        field :num_attention_heads, default: nil
+        field :num_key_value_heads, default: nil
+
+        def initialize(**kwargs)
+          super
+          if @num_attention_heads.nil? && !@hidden_size.nil? && !@head_dim.nil? && @head_dim.to_i > 0
+            @num_attention_heads = @hidden_size / @head_dim
+          end
+          @num_attention_heads ||= 1
+          @num_key_value_heads ||= @num_attention_heads
+          @block_types ||= Array.new(@num_hidden_layers.to_i, "recurrent")
+        end
+
+        def to_recurrent_gemma_dict
+          {
+            "model_type" => @model_type,
+            "attention_bias" => false,
+            "conv1d_width" => 3,
+            "hidden_size" => @hidden_size,
+            "intermediate_size" => @intermediate_size,
+            "logits_soft_cap" => nil,
+            "num_attention_heads" => @num_attention_heads,
+            "num_hidden_layers" => @num_hidden_layers,
+            "num_key_value_heads" => @num_key_value_heads,
+            "rms_norm_eps" => @norm_eps,
+            "rope_theta" => @rope_theta,
+            "attention_window_size" => @attention_window_size,
+            "vocab_size" => @vocab_size,
+            "embeddings_scale_by_sqrt_dim" => false,
+            "block_types" => @block_types,
+          }
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = RecurrentGemma::Model.new(
+            RecurrentGemma::ModelArgs.from_dict(args.to_recurrent_gemma_dict)
+          )
+        end
+
+        def call(inputs, cache: nil)
+          wrapped_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          remapped = {}
+          weights.each do |key, value|
+            remapped[_remap_weight_key(key)] = value
+          end
+          wrapped_model.sanitize(remapped)
+        end
+
+        def layers
+          wrapped_model.layers
+        end
+
+        def make_cache
+          return wrapped_model.make_cache if wrapped_model.respond_to?(:make_cache)
+
+          nil
+        end
+
+        private
+
+        def _remap_weight_key(key)
+          mapped = key.dup
+          mapped = mapped.gsub(/\Ablocks\./, "model.layers.")
+          mapped = mapped.gsub(".time_mix.", ".temporal_block.")
+          mapped
+        end
+      end
+
+      Models.register("rwkv7", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/seed_oss.rb

@@ -0,0 +1,167 @@
+module MlxLm
+  module Models
+    module SeedOSS
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "seed_oss"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 32
+        field :intermediate_size, default: 11008
+        field :num_attention_heads, default: 32
+        field :rms_norm_eps, default: 1e-6
+        field :vocab_size, default: 151936
+        field :num_key_value_heads, default: nil
+        field :head_dim, default: nil
+        field :max_position_embeddings, default: nil
+        field :attention_bias, default: false
+        field :attention_out_bias, default: false
+        field :mlp_bias, default: false
+        field :rope_theta, default: 10000.0
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        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)
+
+          input_bias = args.attention_bias
+          output_bias = args.attention_out_bias
+
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: input_bias)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: input_bias)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: input_bias)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: output_bias)
+
+          self.rope = MlxLm::Models.initialize_rope(
+            @head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            max_position_embeddings: args.max_position_embeddings
+          )
+        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(dim, hidden_dim, bias: false)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: bias)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: bias)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: bias)
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(gate_proj.call(x), up_proj.call(x)))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args.hidden_size, args.intermediate_size, bias: args.mlp_bias)
+          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 SeedModel < 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
+          self.model_type = args.model_type
+          self.model = SeedModel.new(args)
+          self.tie_word_embeddings = args.tie_word_embeddings
+          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)
+          h = model.call(inputs, cache: cache)
+          if tie_word_embeddings
+            model.embed_tokens.as_linear(h)
+          else
+            lm_head.call(h)
+          end
+        end
+
+        def sanitize(weights)
+          weights.delete("lm_head.weight") if tie_word_embeddings
+          weights
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("seed_oss", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/smollm3.rb

@@ -0,0 +1,89 @@
+module MlxLm
+  module Models
+    module SmolLM3
+      class ModelArgs < Llama::ModelArgs
+        field :model_type, default: "smollm3"
+        field :no_rope_layer_interval, default: 4
+        field :no_rope_layers, default: nil
+
+        def initialize(**kwargs)
+          super
+
+          if @no_rope_layers.nil?
+            @no_rope_layers = Array.new(@num_hidden_layers) do |i|
+              ((i + 1) % @no_rope_layer_interval).zero? ? 0 : 1
+            end
+          elsif @no_rope_layers.length != @num_hidden_layers
+            raise ArgumentError, "`no_rope_layers` length mismatch"
+          end
+        end
+      end
+
+      class NoPE < MLX::NN::Module
+        def call(x, offset: 0)
+          x
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.model = Llama::LlamaModel.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+
+          args.no_rope_layers.each_with_index do |use_rope, idx|
+            next if use_rope && use_rope != 0
+
+            model.layers[idx].self_attn.rope = NoPE.new
+          end
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          out = if input_embeddings.nil?
+            model.call(inputs, cache: cache)
+          else
+            _call_with_input_embeddings(input_embeddings, cache)
+          end
+
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def layers
+          model.layers
+        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
+
+        private
+
+        def _call_with_input_embeddings(input_embeddings, cache)
+          h = input_embeddings
+          layer_cache = cache || [nil] * model.layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          model.layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          model.norm.call(h)
+        end
+      end
+
+      Models.register("smollm3", Model, ModelArgs)
+    end
+  end
+end