mlx 0.30.7 → 0.30.7.3

data/lib/mlx/dsl/generate.rb

@@ -0,0 +1,193 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Generate
+      def initialize(
+        model:,
+        tokenizer: nil,
+        eos_id: nil,
+        sampler: nil,
+        mode: :decoder_only,
+        decoder_start_id: nil
+      )
+        @model = model
+        @tokenizer = tokenizer
+        @eos_id = eos_id
+        @sampler = { strategy: :argmax }.merge((sampler || {}).transform_keys(&:to_sym))
+        @mode = mode.to_sym
+        @decoder_start_id = decoder_start_id
+      end
+
+      def each_token(prompt: nil, input_ids: nil, max_tokens: 128, **kwargs)
+        return enum_for(__method__, prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) unless block_given?
+
+        case @mode
+        when :decoder_only
+          __dsl_each_decoder_only(prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) do |id, chunk|
+            yield id, chunk
+          end
+        when :encoder_decoder
+          __dsl_each_encoder_decoder(prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) do |id, chunk|
+            yield id, chunk
+          end
+        else
+          raise ArgumentError, "unsupported generation mode: #{@mode.inspect}"
+        end
+
+        self
+      end
+
+      private
+
+      def __dsl_each_decoder_only(prompt:, input_ids:, max_tokens:, **kwargs)
+        tokens = input_ids.nil? ? __dsl_encode(prompt) : input_ids
+        model_input = __dsl_input_array(tokens)
+        logits, cache = __dsl_decode_step(model_input, cache: nil, **kwargs)
+
+        max_tokens.to_i.times do
+          token = __dsl_sample(__dsl_last_logits(logits))
+          token_id = __dsl_token_id(token)
+          chunk = __dsl_decode_token(token_id)
+          yield token_id, chunk
+          break if !@eos_id.nil? && token_id == @eos_id
+
+          next_input = MLX::Core.array([[token_id]], MLX::Core.int32)
+          logits, cache = __dsl_decode_step(next_input, cache: cache, **kwargs)
+        end
+      end
+
+      def __dsl_each_encoder_decoder(prompt:, input_ids:, max_tokens:, **kwargs)
+        tokens = input_ids.nil? ? __dsl_encode(prompt) : input_ids
+        source = __dsl_input_array(tokens)
+
+        if @model.respond_to?(:encode) && @model.respond_to?(:decode)
+          memory = @model.encode(source)
+          start_id = __dsl_decoder_start_id
+          decoder_input = MLX::Core.array([[start_id]], MLX::Core.int32)
+          cache = nil
+
+          max_tokens.to_i.times do
+            decoded = @model.decode(decoder_input, memory, cache: cache, **kwargs)
+            logits, cache = __dsl_split_logits_and_cache(decoded, cache)
+            token = __dsl_sample(__dsl_last_logits(logits))
+            token_id = __dsl_token_id(token)
+            chunk = __dsl_decode_token(token_id)
+            yield token_id, chunk
+            break if !@eos_id.nil? && token_id == @eos_id
+
+            decoder_input = MLX::Core.array([[token_id]], MLX::Core.int32)
+          end
+          return
+        end
+
+        # Fallback path for model.call-style APIs.
+        __dsl_each_decoder_only(prompt: prompt, input_ids: tokens, max_tokens: max_tokens, **kwargs) do |id, chunk|
+          yield id, chunk
+        end
+      end
+
+      def __dsl_decode_step(input_ids, cache:, **kwargs)
+        output = @model.call(input_ids, cache: cache, **kwargs)
+        __dsl_split_logits_and_cache(output, cache)
+      end
+
+      def __dsl_split_logits_and_cache(output, fallback_cache)
+        if output.is_a?(Array) && output.length == 2
+          [output[0], output[1]]
+        else
+          [output, fallback_cache]
+        end
+      end
+
+      def __dsl_last_logits(logits)
+        return logits if logits.ndim == 2
+        return logits if logits.ndim == 1
+
+        index = MLX::Core.array([logits.shape[1] - 1], MLX::Core.int32)
+        MLX::Core.squeeze(MLX::Core.take(logits, index, 1), 1)
+      end
+
+      def __dsl_sample(logits)
+        strategy = @sampler.fetch(:strategy, :argmax).to_sym
+        temperature = @sampler.fetch(:temperature, 1.0).to_f
+
+        return MLX::Core.argmax(logits, -1) if strategy == :argmax || temperature.zero?
+
+        case strategy
+        when :top_k
+          __dsl_top_k_sample(logits, k: Integer(@sampler.fetch(:k, 40)), temperature: temperature)
+        when :temperature, :categorical
+          __dsl_temperature_sample(logits, temperature: temperature)
+        else
+          raise ArgumentError, "unsupported sampler strategy: #{strategy.inspect}"
+        end
+      end
+
+      def __dsl_temperature_sample(logits, temperature:)
+        scaled = if temperature == 1.0
+          logits
+        else
+          MLX::Core.multiply(logits, 1.0 / temperature)
+        end
+        MLX::Core.categorical(scaled)
+      end
+
+      def __dsl_top_k_sample(logits, k:, temperature:)
+        rows = logits.ndim == 1 ? [logits.to_a] : logits.to_a
+        masked = rows.map do |row|
+          pairs = row.each_with_index.sort_by { |(value, _index)| -value }
+          keep = pairs.first([k, row.length].min).map(&:last)
+          filtered = Array.new(row.length, -Float::INFINITY)
+          keep.each { |idx| filtered[idx] = row[idx] }
+          filtered
+        end
+        masked_logits = MLX::Core.array(masked, logits.dtype)
+        __dsl_temperature_sample(masked_logits, temperature: temperature)
+      end
+
+      def __dsl_encode(prompt)
+        raise ArgumentError, "prompt/input_ids required when tokenizer is unavailable" if @tokenizer.nil?
+
+        @tokenizer.encode(prompt.to_s)
+      end
+
+      def __dsl_input_array(tokens)
+        if tokens.is_a?(MLX::Core::Array)
+          return tokens if tokens.ndim > 1
+
+          return MLX::Core.expand_dims(tokens.astype(MLX::Core.int32), 0)
+        end
+
+        arr = tokens.to_a
+        nested = arr.empty? ? [[]] : (arr.first.is_a?(Array) ? arr : [arr])
+        MLX::Core.array(nested, MLX::Core.int32)
+      end
+
+      def __dsl_token_id(token)
+        return token.item.to_i if token.respond_to?(:item)
+
+        value = token.to_a
+        if value.is_a?(Array)
+          first = value.first
+          return first.first.to_i if first.is_a?(Array)
+          return first.to_i
+        end
+        value.to_i
+      end
+
+      def __dsl_decode_token(token_id)
+        return nil if @tokenizer.nil? || !@tokenizer.respond_to?(:decode)
+
+        @tokenizer.decode([token_id])
+      end
+
+      def __dsl_decoder_start_id
+        return @decoder_start_id unless @decoder_start_id.nil?
+        return @tokenizer.decoder_start_id if !@tokenizer.nil? && @tokenizer.respond_to?(:decoder_start_id)
+
+        raise ArgumentError, "decoder_start_id is required for encoder-decoder mode"
+      end
+    end
+  end
+end

data/lib/mlx/dsl/graph_modules.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Callable < MLX::NN::Module
+      def initialize(callable = nil, &block)
+        super()
+        if !callable.nil? && block_given?
+          raise ArgumentError, "callable layer accepts either a callable argument or block, not both"
+        end
+
+        @callable = callable.nil? ? block : callable
+        unless @callable.respond_to?(:call)
+          raise ArgumentError, "callable layer requires a callable argument or block"
+        end
+      end
+
+      def call(*args, **kwargs)
+        return @callable.call(*args) if kwargs.empty?
+
+        @callable.call(*args, **kwargs)
+      end
+    end
+
+    class Residual < MLX::NN::Module
+      def initialize(module_obj)
+        super()
+        self.module_obj = module_obj
+      end
+
+      def call(*args, **kwargs)
+        raise ArgumentError, "residual module expects at least one positional input" if args.empty?
+
+        identity = args[0]
+        transformed = module_obj.call(*args, **kwargs)
+        MLX::Core.add(identity, transformed)
+      end
+    end
+
+    class Parallel < MLX::NN::Module
+      def initialize(*modules)
+        super()
+        self.layers = modules
+      end
+
+      def call(*args, **kwargs)
+        layers.map do |layer|
+          layer.call(*args, **kwargs)
+        end
+      end
+    end
+
+    class Concat < MLX::NN::Module
+      def initialize(*modules, axis: -1)
+        super()
+        self.layers = modules
+        @axis = axis
+      end
+
+      def call(*args, **kwargs)
+        outputs = layers.map do |layer|
+          layer.call(*args, **kwargs)
+        end
+        MLX::Core.concatenate(outputs, @axis)
+      end
+    end
+
+    class Reduce < MLX::NN::Module
+      def initialize(*modules, mode: :sum)
+        super()
+        self.layers = modules
+        @mode = mode.to_sym
+      end
+
+      def call(*args, **kwargs)
+        outputs = layers.map do |layer|
+          layer.call(*args, **kwargs)
+        end
+
+        case @mode
+        when :sum
+          outputs.reduce do |acc, item|
+            MLX::Core.add(acc, item)
+          end
+        else
+          raise ArgumentError, "unsupported reduce mode: #{@mode.inspect}"
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl/kv_cache.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class KVCache
+      attr_reader :num_layers
+
+      def initialize(num_layers:)
+        @num_layers = Integer(num_layers)
+        raise ArgumentError, "num_layers must be non-negative" if @num_layers.negative?
+
+        @layers = Array.new(@num_layers)
+      end
+
+      def layer(index)
+        @layers.fetch(__dsl_index(index))
+      end
+
+      def []=(index, value)
+        @layers[__dsl_index(index)] = value
+      end
+
+      def offset(layer:)
+        state = self.layer(layer)
+        return 0 if state.nil?
+
+        keys, = state
+        keys.shape[2]
+      end
+
+      def append(layer:, keys:, values:)
+        idx = __dsl_index(layer)
+        current = @layers[idx]
+        if current.nil?
+          @layers[idx] = [keys, values]
+          return @layers[idx]
+        end
+
+        key_cache, value_cache = current
+        next_keys = MLX::Core.concatenate([key_cache, keys], 2)
+        next_values = MLX::Core.concatenate([value_cache, values], 2)
+        @layers[idx] = [next_keys, next_values]
+      end
+
+      def truncate!(tokens:, layer: nil)
+        keep = Integer(tokens)
+        if layer.nil?
+          @layers.each_index { |idx| __dsl_truncate_layer!(idx, keep) }
+        else
+          __dsl_truncate_layer!(__dsl_index(layer), keep)
+        end
+        self
+      end
+
+      def reset!(layer: nil)
+        if layer.nil?
+          @layers.map! { nil }
+        else
+          @layers[__dsl_index(layer)] = nil
+        end
+        self
+      end
+
+      private
+
+      def __dsl_index(index)
+        idx = Integer(index)
+        if idx.negative? || idx >= @num_layers
+          raise IndexError, "layer index #{idx} out of range (0...#{@num_layers})"
+        end
+
+        idx
+      end
+
+      def __dsl_truncate_layer!(idx, keep)
+        state = @layers[idx]
+        return if state.nil?
+
+        if keep <= 0
+          @layers[idx] = nil
+          return
+        end
+
+        keys, values = state
+        total = keys.shape[2]
+        return if keep >= total
+
+        start = total - keep
+        indices = MLX::Core.arange(start, total, 1, MLX::Core.int32)
+        trimmed_keys = MLX::Core.take(keys, indices, 2)
+        trimmed_values = MLX::Core.take(values, indices, 2)
+        @layers[idx] = [trimmed_keys, trimmed_values]
+      end
+    end
+  end
+end

data/lib/mlx/dsl/masks.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Masks
+      module_function
+
+      def causal(length:, offset: 0, dtype: MLX::Core.float32)
+        length = Integer(length)
+        offset = Integer(offset)
+        raise ArgumentError, "length must be non-negative" if length.negative?
+        raise ArgumentError, "offset must be non-negative" if offset.negative?
+
+        rinds = MLX::Core.arange(0, offset + length, 1)
+        linds = if offset.zero?
+          rinds
+        else
+          MLX::Core.arange(offset, offset + length, 1)
+        end
+        lhs = MLX::Core.expand_dims(linds, 1)
+        rhs = MLX::Core.expand_dims(rinds, 0)
+        mask = MLX::Core.less(lhs, rhs).astype(dtype)
+        min_value = if MLX::Core.respond_to?(:finfo)
+          MLX::Core.finfo(dtype).min
+        else
+          -1e9
+        end
+        MLX::Core.multiply(mask, min_value)
+      end
+    end
+  end
+end

data/lib/mlx/dsl/model.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Model < MLX::NN::Module
+      include ModelMixin
+    end
+  end
+end