mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/load_utils.rb

@@ -0,0 +1,87 @@
+require "json"
+
+module MlxLm
+  module LoadUtils
+    module_function
+
+    # Load a model and tokenizer from a local directory.
+    #
+    # @param model_path [String] Path to the model directory
+    # @param tokenizer_config [Hash] Additional tokenizer config overrides
+    # @return [Array(nn::Module, TokenizerWrapper)] The loaded model and tokenizer
+    def load(model_path, tokenizer_config: nil)
+      model, config = load_model(model_path)
+      tokenizer = load_tokenizer(model_path)
+      [model, tokenizer]
+    end
+
+    # Load model from a local directory containing config.json and safetensors.
+    #
+    # @param model_path [String] Path to the model directory
+    # @return [Array(nn::Module, Hash)] The loaded model and config
+    def load_model(model_path)
+      config = Config.load(model_path)
+
+      # Get model and args classes from registry
+      model_class, args_class = Models.get_classes(config)
+
+      # Instantiate model args from config
+      model_args = args_class.from_dict(config)
+
+      # Create model
+      model = model_class.new(model_args)
+
+      # Load weights
+      weights = WeightUtils.load_sharded_safetensors(model_path)
+
+      # Apply model-specific weight sanitization
+      if model.respond_to?(:sanitize)
+        weights = model.sanitize(weights)
+      end
+
+      # Apply quantization if config specifies it
+      quantization = config["quantization"]
+      if quantization
+        group_size = quantization["group_size"] || 64
+        bits = quantization["bits"] || 4
+        Quantize.quantize_model(model, group_size: group_size, bits: bits, weights: weights)
+      end
+
+      # Load weights into model
+      model.load_weights(weights, strict: false)
+
+      [model, config]
+    end
+
+    # Load tokenizer from a local directory.
+    #
+    # @param model_path [String] Path containing tokenizer files
+    # @return [TokenizerWrapper] The loaded tokenizer
+    def load_tokenizer(model_path)
+      tokenizer_path = File.join(model_path, "tokenizer.json")
+      raise "Tokenizer not found at #{tokenizer_path}" unless File.exist?(tokenizer_path)
+
+      tokenizer = Tokenizers::Tokenizer.from_file(tokenizer_path)
+
+      # Try to load tokenizer config for EOS token
+      config_path = File.join(model_path, "tokenizer_config.json")
+      eos_token = nil
+      eos_token_id = nil
+      if File.exist?(config_path)
+        tc = JSON.parse(File.read(config_path))
+        eos_token = tc["eos_token"]
+        eos_token = eos_token["content"] if eos_token.is_a?(Hash)
+      end
+
+      # Try to get eos_token_id from config.json
+      model_config_path = File.join(model_path, "config.json")
+      if File.exist?(model_config_path)
+        mc = JSON.parse(File.read(model_config_path))
+        eos_token_id = mc["eos_token_id"]
+        eos_token_id = eos_token_id.is_a?(::Array) ? eos_token_id : [eos_token_id].compact
+      end
+
+      TokenizerWrapper.new(tokenizer, eos_token: eos_token, eos_token_id: eos_token_id)
+    end
+  end
+end

data/lib/mlx_lm/model_args.rb

@@ -0,0 +1,54 @@
+module MlxLm
+  # Base class for model configuration arguments.
+  # Mirrors Python's BaseModelArgs dataclass with from_dict filtering.
+  class BaseModelArgs
+    def self.fields
+      @fields ||= {}
+    end
+
+    def self.field(name, default: :__required__)
+      fields[name] = default
+
+      attr_accessor name
+
+      if default == :__required__
+        # no default
+      else
+        define_method(:"default_#{name}") { default }
+      end
+    end
+
+    def self.inherited(subclass)
+      super
+      # Copy parent fields into subclass
+      subclass.instance_variable_set(:@fields, fields.dup)
+    end
+
+    def self.from_dict(params)
+      known = {}
+      fields.each do |name, default|
+        str_name = name.to_s
+        if params.key?(str_name)
+          known[name] = params[str_name]
+        elsif default != :__required__
+          known[name] = default
+        end
+      end
+      new(**known)
+    end
+
+    def initialize(**kwargs)
+      kwargs.each do |k, v|
+        if self.class.fields.key?(k)
+          instance_variable_set(:"@#{k}", v)
+        end
+      end
+      # Set defaults for any fields not provided
+      self.class.fields.each do |name, default|
+        next if kwargs.key?(name)
+        next if default == :__required__
+        instance_variable_set(:"@#{name}", default)
+      end
+    end
+  end
+end

data/lib/mlx_lm/models/activations.rb

@@ -0,0 +1,46 @@
+module MlxLm
+  module Models
+    module Activations
+      module_function
+
+      def swiglu(gate, x)
+        MLX::NN.silu(gate) * x
+      end
+
+      def xielu(x, alpha_p, alpha_n, beta, eps)
+        mx = MLX::Core
+        alpha_p = MLX::NN.softplus(alpha_p)
+        alpha_n = beta + MLX::NN.softplus(alpha_n)
+
+        mx.where(
+          mx.greater(x, 0.0),
+          alpha_p * mx.square(x) + beta * x,
+          (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x
+        )
+      end
+
+      class XieLU < MLX::NN::Module
+        def initialize(
+          alpha_p_init: 0.8,
+          alpha_n_init: 0.8,
+          beta: 0.5,
+          eps: -1e-6
+        )
+          super()
+          mx = MLX::Core
+          alpha_p_tensor = mx.array(alpha_p_init)
+          alpha_n_tensor = mx.array(alpha_n_init - beta)
+
+          self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1.0)
+          self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1.0)
+          self.beta = mx.array(beta)
+          self.eps = mx.array(eps)
+        end
+
+        def call(x)
+          Activations.xielu(x, alpha_p, alpha_n, beta, eps)
+        end
+      end
+    end
+  end
+end

data/lib/mlx_lm/models/afm7.rb

@@ -0,0 +1,131 @@
+require_relative "afmoe"
+
+module MlxLm
+  module Models
+    module Afm7
+      class ModelArgs < Afmoe::ModelArgs
+        field :model_type, default: "afm7"
+        field :hidden_dim, default: nil
+        field :num_layers, default: nil
+        field :num_kv_reuse_layers, default: 0
+        field :num_heads, default: nil
+        field :num_kv_heads, default: nil
+        field :hidden_dim_scale_factor, default: nil
+
+        def initialize(**kwargs)
+          afm7_style = _afm7_style_kwargs?(kwargs)
+          super
+
+          @hidden_size = @hidden_dim if kwargs.key?(:hidden_dim) && !@hidden_dim.nil?
+          @num_hidden_layers = @num_layers if kwargs.key?(:num_layers) && !@num_layers.nil?
+          @num_attention_heads = @num_heads if kwargs.key?(:num_heads) && !@num_heads.nil?
+          @num_key_value_heads = @num_kv_heads if kwargs.key?(:num_kv_heads) && !@num_kv_heads.nil?
+
+          if kwargs.key?(:hidden_dim_scale_factor) && !@hidden_dim_scale_factor.nil? && !@hidden_size.nil?
+            @intermediate_size = (@hidden_size * @hidden_dim_scale_factor.to_f).to_i
+          end
+
+          if !@hidden_size.nil? && !@num_attention_heads.nil? && @num_attention_heads.to_i > 0
+            @head_dim = @hidden_size / @num_attention_heads
+          end
+
+          if kwargs.key?(:num_kv_reuse_layers) && !@num_hidden_layers.nil?
+            @num_dense_layers = [@num_hidden_layers.to_i - @num_kv_reuse_layers.to_i, 0].max
+          elsif afm7_style && !@num_hidden_layers.nil?
+            @num_dense_layers = @num_hidden_layers
+          end
+
+          if afm7_style
+            @num_experts = 1 unless kwargs.key?(:num_experts)
+            @num_experts_per_tok = 1 unless kwargs.key?(:num_experts_per_tok)
+            @num_shared_experts = 0 unless kwargs.key?(:num_shared_experts)
+            @mup_enabled = false unless kwargs.key?(:mup_enabled)
+            @layer_types = Array.new(@num_hidden_layers) { "full_attention" } unless kwargs.key?(:layer_types)
+          end
+
+          @num_key_value_heads ||= @num_attention_heads
+          @layer_types ||= Array.new(@num_hidden_layers) { "full_attention" } unless @num_hidden_layers.nil?
+        end
+
+        def to_afmoe_dict
+          {
+            "model_type" => @model_type,
+            "layer_types" => @layer_types,
+            "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,
+            "head_dim" => @head_dim,
+            "max_position_embeddings" => @max_position_embeddings,
+            "rms_norm_eps" => @rms_norm_eps,
+            "rope_theta" => @rope_theta,
+            "rope_scaling" => @rope_scaling,
+            "tie_word_embeddings" => @tie_word_embeddings,
+            "num_experts" => @num_experts,
+            "num_experts_per_tok" => @num_experts_per_tok,
+            "num_shared_experts" => @num_shared_experts,
+            "num_dense_layers" => @num_dense_layers,
+            "route_norm" => @route_norm,
+            "route_scale" => @route_scale,
+            "score_func" => @score_func,
+            "n_group" => @n_group,
+            "topk_group" => @topk_group,
+            "sliding_window" => @sliding_window,
+            "mup_enabled" => @mup_enabled,
+          }
+        end
+
+        private
+
+        def _afm7_style_kwargs?(kwargs)
+          kwargs.key?(:hidden_dim) ||
+            kwargs.key?(:num_layers) ||
+            kwargs.key?(:num_heads) ||
+            kwargs.key?(:num_kv_heads) ||
+            kwargs.key?(:num_kv_reuse_layers) ||
+            kwargs.key?(:hidden_dim_scale_factor)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = Afmoe::Model.new(Afmoe::ModelArgs.from_dict(args.to_afmoe_dict))
+        end
+
+        def call(inputs, cache: nil)
+          wrapped_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          wrapped_model.sanitize(weights)
+        end
+
+        def layers
+          wrapped_model.layers
+        end
+
+        def make_cache
+          return nil unless wrapped_model.respond_to?(:make_cache)
+
+          wrapped_model.make_cache
+        end
+
+        def cast_predicate
+          wrapped_model.cast_predicate
+        end
+
+        def quant_predicate
+          wrapped_model.quant_predicate
+        end
+      end
+
+      Models.register("afm7", Model, ModelArgs)
+    end
+  end
+end