mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/llama4_text.rb

@@ -0,0 +1,195 @@
+module MlxLm
+  module Models
+    module Llama4Text
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "llama4_text"
+        field :hidden_size
+        field :num_attention_heads
+        field :num_hidden_layers
+        field :vocab_size
+        field :intermediate_size, default: nil
+        field :intermediate_size_mlp, default: nil
+        field :num_key_value_heads, default: nil
+        field :rms_norm_eps, default: 1e-5
+        field :rope_theta, default: 10_000.0
+        field :head_dim, default: nil
+        field :tie_word_embeddings, default: true
+        field :no_rope_layers, default: nil
+        field :use_qk_norm, default: false
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+          @intermediate_size_mlp ||= @intermediate_size
+
+          if @no_rope_layers.nil?
+            @no_rope_layers = Array.new(@num_hidden_layers, 1)
+          elsif @no_rope_layers.length != @num_hidden_layers
+            raise ArgumentError, "`no_rope_layers` length mismatch"
+          end
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args, use_rope)
+          super()
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = @head_dim**(-0.5)
+          @use_rope = !!use_rope
+          @use_qk_norm = !!args.use_qk_norm
+          @rms_norm_eps = args.rms_norm_eps
+
+          self.q_proj = MLX::NN::Linear.new(args.hidden_size, @n_heads * @head_dim, bias: false)
+          self.k_proj = MLX::NN::Linear.new(args.hidden_size, @n_kv_heads * @head_dim, bias: false)
+          self.v_proj = MLX::NN::Linear.new(args.hidden_size, @n_kv_heads * @head_dim, bias: false)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, args.hidden_size, bias: false)
+
+          if @use_rope
+            self.rope = MLX::NN::RoPE.new(@head_dim, traditional: true, base: args.rope_theta)
+          end
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x)
+          keys = k_proj.call(x)
+          values = v_proj.call(x)
+
+          queries = queries.reshape([b, l, @n_heads, @head_dim])
+          keys = keys.reshape([b, l, @n_kv_heads, @head_dim])
+
+          if @use_qk_norm
+            queries = mx.rms_norm(queries, nil, @rms_norm_eps)
+            keys = mx.rms_norm(keys, nil, @rms_norm_eps)
+          end
+
+          queries = queries.transpose([0, 2, 1, 3])
+          keys = keys.transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if @use_rope
+            if cache
+              queries = rope.call(queries, offset: cache.offset)
+              keys = rope.call(keys, offset: cache.offset)
+            else
+              queries = rope.call(queries)
+              keys = rope.call(keys)
+            end
+          end
+
+          keys, values = cache.update_and_fetch(keys, values) if cache
+
+          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, intermediate_size)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(dim, intermediate_size, bias: false)
+          self.up_proj = MLX::NN::Linear.new(dim, intermediate_size, bias: false)
+          self.down_proj = MLX::NN::Linear.new(intermediate_size, dim, bias: false)
+        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, use_rope)
+          super()
+          self.self_attn = Attention.new(args, use_rope)
+          self.feed_forward = MLP.new(args.hidden_size, args.intermediate_size_mlp)
+          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 = feed_forward.call(post_attention_layernorm.call(h))
+          h + r
+        end
+      end
+
+      class LanguageModel < 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) do |i|
+            TransformerBlock.new(args, args.no_rope_layers[i])
+          end
+          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 = _create_attention_mask(h, layer_cache[0])
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache)
+          return cache.make_mask(h.shape[1]) if cache && cache.respond_to?(:make_mask)
+          return nil if h.shape[1] == 1
+
+          "causal"
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.model = LanguageModel.new(args)
+          self.output = nil
+          unless args.tie_word_embeddings
+            self.output = 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 @args.tie_word_embeddings
+            model.embed_tokens.as_linear(h)
+          else
+            output.call(h)
+          end
+        end
+
+        def sanitize(weights)
+          sanitized = weights.reject do |k, _|
+            k.include?("self_attn.rotary_emb.inv_freq") || k.include?("self_attn.rope.inv_freq")
+          end
+          if @args.tie_word_embeddings
+            sanitized.delete("output.weight")
+            sanitized.delete("lm_head.weight")
+          end
+          sanitized
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("llama4_text", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/longcat_flash.rb

@@ -0,0 +1,153 @@
+require_relative "glm4_moe_lite"
+
+module MlxLm
+  module Models
+    module LongcatFlash
+      class ModelArgs < Glm4MoeLite::ModelArgs
+        field :model_type, default: "longcat_flash"
+        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_experts, default: nil
+        field :num_local_experts, default: nil
+        field :num_shared_experts, default: nil
+        field :top_k, default: nil
+        field :score_function, 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" => "n_routed_experts",
+            "num_experts" => "n_routed_experts",
+            "num_shared_experts" => "n_shared_experts",
+            "top_k" => "num_experts_per_tok",
+            "score_function" => "scoring_func",
+          }.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"] ||= "longcat_flash"
+          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?
+          @n_routed_experts = @num_local_experts if kwargs.key?(:num_local_experts) && !kwargs.key?(:n_routed_experts) && !@num_local_experts.nil?
+          @n_routed_experts = @num_experts if kwargs.key?(:num_experts) && !kwargs.key?(:n_routed_experts) && !kwargs.key?(:num_local_experts) && !@num_experts.nil?
+          @n_shared_experts = @num_shared_experts if kwargs.key?(:num_shared_experts) && !kwargs.key?(:n_shared_experts) && !@num_shared_experts.nil?
+          @num_experts_per_tok = @top_k if kwargs.key?(:top_k) && !kwargs.key?(:num_experts_per_tok) && !@top_k.nil?
+          @scoring_func = @score_function if kwargs.key?(:score_function) && !kwargs.key?(:scoring_func) && !@score_function.nil?
+          @num_key_value_heads ||= @num_attention_heads
+        end
+
+        def to_glm4_moe_lite_dict
+          {
+            "model_type" => @model_type,
+            "vocab_size" => @vocab_size,
+            "hidden_size" => @hidden_size,
+            "intermediate_size" => @intermediate_size,
+            "moe_intermediate_size" => @moe_intermediate_size,
+            "num_hidden_layers" => @num_hidden_layers,
+            "num_attention_heads" => @num_attention_heads,
+            "num_key_value_heads" => @num_key_value_heads,
+            "n_shared_experts" => @n_shared_experts,
+            "n_routed_experts" => @n_routed_experts,
+            "routed_scaling_factor" => @routed_scaling_factor,
+            "kv_lora_rank" => @kv_lora_rank,
+            "q_lora_rank" => @q_lora_rank,
+            "qk_rope_head_dim" => @qk_rope_head_dim,
+            "qk_nope_head_dim" => @qk_nope_head_dim,
+            "v_head_dim" => @v_head_dim,
+            "topk_method" => @topk_method,
+            "scoring_func" => @scoring_func,
+            "norm_topk_prob" => @norm_topk_prob,
+            "n_group" => @n_group,
+            "topk_group" => @topk_group,
+            "num_experts_per_tok" => @num_experts_per_tok,
+            "moe_layer_freq" => @moe_layer_freq,
+            "first_k_dense_replace" => @first_k_dense_replace,
+            "max_position_embeddings" => @max_position_embeddings,
+            "rms_norm_eps" => @rms_norm_eps,
+            "rope_theta" => @rope_theta,
+            "rope_scaling" => @rope_scaling,
+            "attention_bias" => @attention_bias,
+            "attention_dropout" => @attention_dropout,
+            "partial_rotary_factor" => @partial_rotary_factor,
+            "tie_word_embeddings" => @tie_word_embeddings,
+            "num_nextn_predict_layers" => @num_nextn_predict_layers,
+            "quantization" => @quantization,
+          }
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = Glm4MoeLite::Model.new(
+            Glm4MoeLite::ModelArgs.from_dict(args.to_glm4_moe_lite_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
+
+        private
+
+        def _remap_weight_key(key)
+          mapped = key.dup
+          mapped = mapped.gsub(".attention.", ".self_attn.")
+          mapped = mapped.gsub(".block_sparse_moe.", ".mlp.")
+          mapped = mapped.gsub(".mlp.router.", ".mlp.gate.")
+          mapped
+        end
+      end
+
+      Models.register("longcat_flash", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/longcat_flash_ngram.rb

@@ -0,0 +1,137 @@
+require_relative "longcat_flash"
+
+module MlxLm
+  module Models
+    module LongcatFlashNgram
+      class ModelArgs < LongcatFlash::ModelArgs
+        field :model_type, default: "longcat_flash_ngram"
+        field :attention_method, default: nil
+        field :zero_expert_type, default: "identity"
+        field :moe_topk, default: nil
+        field :expert_ffn_hidden_size, default: nil
+        field :zero_expert_num, default: nil
+        field :num_layers, default: nil
+        field :ngram_vocab_size_ratio, default: 78
+        field :emb_neighbor_num, default: 4
+        field :emb_split_num, default: 4
+        field :mla_scale_q_lora, default: nil
+        field :mla_scale_kv_lora, default: nil
+        field :router_bias, default: false
+
+        def self.from_dict(params)
+          normalized = params.each_with_object({}) do |(key, value), out|
+            out[key.to_s] = value
+          end
+
+          {
+            "num_layers" => "num_hidden_layers",
+            "moe_topk" => "num_experts_per_tok",
+            "expert_ffn_hidden_size" => "moe_intermediate_size",
+          }.each do |source_key, target_key|
+            next unless normalized.key?(source_key)
+
+            normalized[target_key] = normalized[source_key] unless normalized.key?(target_key)
+          end
+
+          if normalized.key?("n_routed_experts") && normalized.key?("zero_expert_num") && !normalized.key?("num_local_experts")
+            normalized["num_local_experts"] = normalized["n_routed_experts"].to_i + normalized["zero_expert_num"].to_i
+          end
+
+          if normalized.key?("num_attention_heads") && !normalized.key?("num_key_value_heads") && !normalized.key?("num_kv_heads")
+            normalized["num_key_value_heads"] = normalized["num_attention_heads"]
+          end
+
+          normalized["model_type"] ||= "longcat_flash_ngram"
+          super(normalized)
+        end
+
+        def initialize(**kwargs)
+          super
+          @num_hidden_layers = @num_layers if kwargs.key?(:num_layers) && !kwargs.key?(:num_hidden_layers) && !@num_layers.nil?
+          @num_experts_per_tok = @moe_topk if kwargs.key?(:moe_topk) && !kwargs.key?(:num_experts_per_tok) && !@moe_topk.nil?
+          @moe_intermediate_size = @expert_ffn_hidden_size if kwargs.key?(:expert_ffn_hidden_size) && !kwargs.key?(:moe_intermediate_size) && !@expert_ffn_hidden_size.nil?
+
+          if kwargs.key?(:zero_expert_num) && !@zero_expert_num.nil? && !kwargs.key?(:num_local_experts) && !kwargs.key?(:n_routed_experts) && !@n_routed_experts.nil?
+            @n_routed_experts = @n_routed_experts.to_i + @zero_expert_num.to_i
+          end
+
+          if kwargs.key?(:num_attention_heads) && !kwargs.key?(:num_key_value_heads) && !kwargs.key?(:num_kv_heads)
+            @num_key_value_heads = @num_attention_heads
+          end
+
+          @num_key_value_heads ||= @num_attention_heads
+        end
+
+        def to_longcat_flash_dict
+          routed_experts = @n_routed_experts
+          if !@zero_expert_num.nil? && !routed_experts.nil?
+            routed_experts = routed_experts.to_i + @zero_expert_num.to_i
+          end
+
+          dict = to_glm4_moe_lite_dict
+          dict["model_type"] = @model_type
+          dict["n_routed_experts"] = routed_experts unless routed_experts.nil?
+          dict
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = LongcatFlash::Model.new(
+            LongcatFlash::ModelArgs.from_dict(args.to_longcat_flash_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[_to_longcat_flash_key(key)] = value
+          end
+
+          sanitized = wrapped_model.sanitize(remapped)
+          restored = {}
+          sanitized.each do |key, value|
+            restored[_from_longcat_flash_key(key)] = value
+          end
+          restored
+        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
+
+        private
+
+        def _to_longcat_flash_key(key)
+          key.to_s.gsub("model.ngram_embeddings.word_embeddings.", "model.embed_tokens.")
+        end
+
+        def _from_longcat_flash_key(key)
+          key.to_s.gsub("model.embed_tokens.", "model.ngram_embeddings.word_embeddings.")
+        end
+      end
+
+      Models.register("longcat_flash_ngram", Model, ModelArgs)
+    end
+  end
+end