mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/jamba.rb

@@ -0,0 +1,158 @@
+require_relative "falcon_h1"
+
+module MlxLm
+  module Models
+    module Jamba
+      class ModelArgs < FalconH1::ModelArgs
+        field :model_type, default: "jamba"
+        field :attn_layer_offset, default: 1
+        field :attn_layer_period, default: 2
+        field :expert_layer_offset, default: 1
+        field :expert_layer_period, default: 2
+        field :mamba_d_state, default: nil
+        field :mamba_expand, default: nil
+        field :num_experts, default: 1
+        field :num_experts_per_tok, default: 1
+        field :mamba_dt_rank, default: "auto"
+        field :mamba_proj_bias, default: false
+        field :mamba_conv_bias, default: true
+        field :layers_block_type, default: nil
+
+        def initialize(**kwargs)
+          super
+          @mamba_d_conv ||= 4
+          @num_key_value_heads ||= @num_attention_heads
+          @layers_block_type ||= _default_layers_block_type
+          @num_hidden_layers ||= Array(@layers_block_type).length
+          @block_types ||= _to_block_types
+        end
+
+        def to_falcon_h1_dict
+          hidden_size = @hidden_size
+          attention_heads = @num_attention_heads
+          inferred_head_dim = if !@head_dim.nil?
+            @head_dim
+          elsif !hidden_size.nil? && attention_heads.to_i > 0
+            hidden_size / attention_heads
+          else
+            64
+          end
+
+          {
+            "model_type" => @model_type,
+            "attention_bias" => @attention_bias,
+            "head_dim" => inferred_head_dim,
+            "hidden_size" => hidden_size,
+            "intermediate_size" => @intermediate_size,
+            "max_position_embeddings" => @max_position_embeddings,
+            "mamba_d_conv" => @mamba_d_conv,
+            "num_attention_heads" => attention_heads,
+            "num_hidden_layers" => @num_hidden_layers,
+            "num_key_value_heads" => @num_key_value_heads,
+            "rms_norm_eps" => @rms_norm_eps,
+            "rope_theta" => @rope_theta,
+            "vocab_size" => @vocab_size,
+            "tie_word_embeddings" => @tie_word_embeddings,
+            "attention_window_size" => @attention_window_size,
+            "block_types" => @block_types,
+          }
+        end
+
+        private
+
+        def _default_layers_block_type
+          count = @num_hidden_layers.to_i
+          return nil if count <= 0
+
+          period = @attn_layer_period.to_i
+          offset = @attn_layer_offset.to_i
+          period = 1 if period <= 0
+
+          Array.new(count) do |idx|
+            (idx % period == offset) ? "attention" : "mamba"
+          end
+        end
+
+        def _to_block_types
+          return @block_types if @block_types.is_a?(Array) && !@block_types.empty?
+          return nil unless @layers_block_type.is_a?(Array) && !@layers_block_type.empty?
+
+          @layers_block_type.map { |layer_type| layer_type.to_s == "mamba" ? "recurrent" : "attention" }
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = FalconH1::Model.new(
+            FalconH1::ModelArgs.from_dict(args.to_falcon_h1_dict)
+          )
+        end
+
+        def call(inputs, cache: nil)
+          wrapped_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          normalized = weights.dup
+          _stack_experts!(normalized)
+
+          remapped = {}
+          normalized.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 nil unless wrapped_model.respond_to?(:make_cache)
+
+          wrapped_model.make_cache
+        end
+
+        private
+
+        def _stack_experts!(weights)
+          mx = MLX::Core
+
+          @args.num_hidden_layers.to_i.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}.feed_forward"
+            %w[gate_proj up_proj down_proj].each do |projection|
+              %w[weight bias scales biases].each do |param|
+                pattern = /\A#{Regexp.escape(prefix)}\.experts\.(\d+)\.#{projection}\.#{param}\z/
+                matches = weights.keys.filter_map do |key|
+                  match = pattern.match(key)
+                  next nil unless match
+
+                  [match[1].to_i, key]
+                end
+                next if matches.empty?
+
+                stacked = matches.sort_by(&:first).map do |(_, key)|
+                  weights.delete(key)
+                end
+                weights["#{prefix}.switch_mlp.#{projection}.#{param}"] = mx.stack(stacked)
+              end
+            end
+          end
+        end
+
+        def _remap_weight_key(key)
+          mapped = key.dup
+          mapped = mapped.gsub("model.norm.", "model.final_layernorm.")
+          mapped = mapped.gsub(".mixer.", ".mamba.")
+          mapped = mapped.gsub(".feed_forward.router.", ".feed_forward.gate.")
+          mapped
+        end
+      end
+
+      Models.register("jamba", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/kimi_k25.rb

@@ -0,0 +1,98 @@
+require_relative "deepseek"
+
+module MlxLm
+  module Models
+    module KimiK25
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "kimi_k25"
+        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
+
+          model_type = params["model_type"] || params[:model_type] || "kimi_k25"
+          new(model_type: model_type, text_config: params)
+        end
+
+        def initialize(**kwargs)
+          super
+          @text_config = _stringify_keys(@text_config || {})
+          @text_config["model_type"] ||= "deepseek"
+        end
+
+        private
+
+        def _stringify_keys(hash)
+          hash.each_with_object({}) do |(key, value), out|
+            out[key.to_s] = value
+          end
+        end
+      end
+
+      class Model < MLX::NN::Module
+        MULTIMODAL_PREFIXES = %w[
+          vision_tower
+          vision_model
+          multi_modal_projector
+          mm_projector
+        ].freeze
+
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.language_model = DeepSeek::Model.new(
+            DeepSeek::ModelArgs.from_dict(args.text_config)
+          )
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          language_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          language_weights = {}
+          flat_weights = weights.is_a?(Hash) ? weights : weights.to_h
+
+          flat_weights.each do |key, value|
+            next if _multimodal_key?(key)
+
+            normalized_key = key.start_with?("language_model.") ? key.delete_prefix("language_model.") : key
+            language_weights[normalized_key] = value
+          end
+
+          sanitized_language = if language_model.respond_to?(:sanitize)
+            language_model.sanitize(language_weights)
+          else
+            language_weights
+          end
+
+          sanitized_language.each_with_object({}) do |(key, value), out|
+            out["language_model.#{key}"] = value
+          end
+        end
+
+        def model
+          language_model.model
+        end
+
+        def layers
+          model.layers
+        end
+
+        def cast_predicate
+          lambda { |key| !key.include?("e_score_correction_bias") }
+        end
+
+        private
+
+        def _multimodal_key?(key)
+          MULTIMODAL_PREFIXES.any? { |prefix| key == prefix || key.start_with?("#{prefix}.") }
+        end
+      end
+
+      Models.register("kimi_k25", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/kimi_linear.rb

@@ -0,0 +1,124 @@
+require_relative "bailing_moe_linear"
+
+module MlxLm
+  module Models
+    module KimiLinear
+      class ModelArgs < BailingMoeLinear::ModelArgs
+        field :model_type, default: "kimi_linear"
+        field :hidden_dim, default: nil
+        field :ffn_hidden_size, default: nil
+        field :num_layers, default: nil
+        field :num_heads, default: nil
+        field :num_kv_heads, default: nil
+        field :num_local_experts, default: nil
+        field :n_routed_experts, default: nil
+        field :n_shared_experts, default: nil
+        field :top_k, default: nil
+        field :score_func, default: nil
+
+        def self.from_dict(params)
+          normalized = params.each_with_object({}) do |(key, value), out|
+            out[key.to_s] = value
+          end
+
+          {
+            "hidden_dim" => "hidden_size",
+            "ffn_hidden_size" => "intermediate_size",
+            "num_layers" => "num_hidden_layers",
+            "num_heads" => "num_attention_heads",
+            "num_kv_heads" => "num_key_value_heads",
+            "num_local_experts" => "num_experts",
+            "n_routed_experts" => "num_experts",
+            "n_shared_experts" => "num_shared_experts",
+            "top_k" => "num_experts_per_tok",
+            "score_func" => "score_function",
+          }.each do |source_key, target_key|
+            next unless normalized.key?(source_key)
+
+            normalized[target_key] = normalized[source_key] unless normalized.key?(target_key)
+          end
+
+          normalized["model_type"] ||= "kimi_linear"
+          super(normalized)
+        end
+
+        def initialize(**kwargs)
+          super
+          @hidden_size = @hidden_dim if kwargs.key?(:hidden_dim) && !kwargs.key?(:hidden_size) && !@hidden_dim.nil?
+          @intermediate_size = @ffn_hidden_size if kwargs.key?(:ffn_hidden_size) && !kwargs.key?(:intermediate_size) && !@ffn_hidden_size.nil?
+          @num_hidden_layers = @num_layers if kwargs.key?(:num_layers) && !kwargs.key?(:num_hidden_layers) && !@num_layers.nil?
+          @num_attention_heads = @num_heads if kwargs.key?(:num_heads) && !kwargs.key?(:num_attention_heads) && !@num_heads.nil?
+          @num_key_value_heads = @num_kv_heads if kwargs.key?(:num_kv_heads) && !kwargs.key?(:num_key_value_heads) && !@num_kv_heads.nil?
+          @num_experts = @num_local_experts if kwargs.key?(:num_local_experts) && !kwargs.key?(:num_experts) && !@num_local_experts.nil?
+          @num_experts = @n_routed_experts if kwargs.key?(:n_routed_experts) && !kwargs.key?(:num_experts) && !kwargs.key?(:num_local_experts) && !@n_routed_experts.nil?
+          @num_shared_experts = @n_shared_experts if kwargs.key?(:n_shared_experts) && !kwargs.key?(:num_shared_experts) && !@n_shared_experts.nil?
+          @num_experts_per_tok = @top_k if kwargs.key?(:top_k) && !kwargs.key?(:num_experts_per_tok) && !@top_k.nil?
+          @score_function = @score_func if kwargs.key?(:score_func) && !kwargs.key?(:score_function) && !@score_func.nil?
+          @num_key_value_heads ||= @num_attention_heads
+        end
+
+        def to_bailing_moe_linear_dict
+          to_bailing_moe_dict
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = BailingMoeLinear::Model.new(
+            BailingMoeLinear::ModelArgs.from_dict(args.to_bailing_moe_linear_dict)
+          )
+        end
+
+        def call(inputs, cache: nil)
+          wrapped_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          remapped = {}
+          flat_weights = weights.is_a?(Hash) ? weights : weights.to_h
+          flat_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
+
+        def cast_predicate
+          return wrapped_model.cast_predicate if wrapped_model.respond_to?(:cast_predicate)
+
+          lambda { |_key| true }
+        end
+
+        def quant_predicate
+          return wrapped_model.quant_predicate if wrapped_model.respond_to?(:quant_predicate)
+
+          lambda { |_key, _value| true }
+        end
+
+        private
+
+        def _remap_weight_key(key)
+          mapped = key.dup
+          mapped = mapped.gsub(".mlp.router.", ".mlp.gate.")
+          mapped = mapped.gsub("model.embed_tokens.", "model.word_embeddings.")
+          mapped = mapped.gsub("model.tok_embeddings.", "model.word_embeddings.")
+          mapped
+        end
+      end
+
+      Models.register("kimi_linear", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/kimi_vl.rb

@@ -0,0 +1,93 @@
+require_relative "deepseek"
+
+module MlxLm
+  module Models
+    module KimiVL
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "kimi_vl"
+        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
+
+          model_type = params["model_type"] || params[:model_type] || "kimi_vl"
+          new(model_type: model_type, text_config: params)
+        end
+
+        def initialize(**kwargs)
+          super
+          @text_config = _stringify_keys(@text_config || {})
+          @text_config["model_type"] ||= "deepseek"
+        end
+
+        private
+
+        def _stringify_keys(hash)
+          hash.each_with_object({}) do |(key, value), out|
+            out[key.to_s] = value
+          end
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.language_model = DeepSeek::Model.new(
+            DeepSeek::ModelArgs.from_dict(args.text_config)
+          )
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          language_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          language_weights = {}
+          flat_weights = weights.is_a?(Hash) ? weights : weights.to_h
+
+          flat_weights.each do |key, value|
+            next if _drop_key?(key)
+
+            normalized_key = key.start_with?("language_model.") ? key.delete_prefix("language_model.") : key
+            language_weights[normalized_key] = value
+          end
+
+          sanitized_language = if language_model.respond_to?(:sanitize)
+            language_model.sanitize(language_weights)
+          else
+            language_weights
+          end
+
+          sanitized_language.each_with_object({}) do |(key, value), out|
+            out["language_model.#{key}"] = value
+          end
+        end
+
+        def model
+          language_model.model
+        end
+
+        def layers
+          model.layers
+        end
+
+        def cast_predicate
+          lambda { |key| !key.include?("e_score_correction_bias") }
+        end
+
+        private
+
+        def _drop_key?(key)
+          key.include?("vision_tower") ||
+            key.include?("multi_modal_projector") ||
+            key.include?("rotary_emb")
+        end
+      end
+
+      Models.register("kimi_vl", Model, ModelArgs)
+    end
+  end
+end