mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/sample_utils.rb

@@ -0,0 +1,159 @@
+module MlxLm
+  module SampleUtils
+    module_function
+
+    # Build a sampler callable (proc) from the given parameters.
+    # Returns a proc that takes logprobs (mx.array) and returns a token (mx.array).
+    def make_sampler(
+      temp: 0.0,
+      top_p: 0.0,
+      min_p: 0.0,
+      min_tokens_to_keep: 1,
+      top_k: 0
+    )
+      mx = MLX::Core
+
+      if temp == 0
+        return ->(x) { mx.argmax(x, -1) }
+      end
+
+      sampling_methods = []
+      if top_p > 0 && top_p < 1.0
+        sampling_methods << ->(x) { apply_top_p(x, top_p) }
+      end
+      if min_p != 0.0
+        sampling_methods << ->(x) { apply_min_p(x, min_p, min_tokens_to_keep) }
+      end
+      if top_k > 0
+        sampling_methods << ->(x) { apply_top_k(x, top_k) }
+      end
+
+      ->(logprobs) {
+        sampling_methods.each { |method| logprobs = method.call(logprobs) }
+        categorical_sampling(logprobs, temp)
+      }
+    end
+
+    def make_logits_processors(repetition_penalty: nil, repetition_context_size: 20)
+      processors = []
+      if repetition_penalty && repetition_penalty != 0.0
+        processors << make_repetition_penalty(repetition_penalty, repetition_context_size)
+      end
+      processors
+    end
+
+    def apply_top_k(logprobs, top_k)
+      mx = MLX::Core
+      vocab_size = logprobs.shape[-1]
+      raise ArgumentError, "top_k must be in (0, #{vocab_size}]" unless top_k.is_a?(Integer) && top_k > 0 && top_k < vocab_size
+
+      neg_logprobs = mx.negative(logprobs)
+      mask_idx = mx.argpartition(neg_logprobs, top_k - 1, -1)
+      # Get indices after top_k (the ones to mask)
+      rest = mx.split(mask_idx, [top_k], -1)[1]
+      neg_inf = mx.array([-Float::INFINITY], dtype: logprobs.dtype)
+      mx.put_along_axis(logprobs, rest, neg_inf, -1)
+    end
+
+    def apply_min_p(logprobs, min_p, min_tokens_to_keep = 1)
+      mx = MLX::Core
+      raise ArgumentError, "min_p must be in [0, 1]" unless min_p >= 0 && min_p <= 1.0
+
+      # Sort indices in decreasing order
+      neg_logprobs = mx.negative(logprobs)
+      sorted_indices = mx.argsort(neg_logprobs, -1)
+      sorted_logprobs = mx.take_along_axis(logprobs, sorted_indices, -1)
+
+      # Top probability
+      top_logprobs = mx.split(sorted_logprobs, [1], -1)[0]
+
+      # Calculate the min_p threshold
+      scaled_min_p = top_logprobs + Math.log(min_p)
+
+      # Mask tokens below threshold
+      tokens_to_remove = mx.less(sorted_logprobs, scaled_min_p)
+
+      neg_inf = mx.array(-Float::INFINITY, dtype: sorted_logprobs.dtype)
+      selected_logprobs = mx.where(tokens_to_remove, neg_inf, sorted_logprobs)
+
+      # Restore the top min_tokens_to_keep tokens regardless
+      if min_tokens_to_keep > 0
+        top_sorted = mx.split(sorted_logprobs, [min_tokens_to_keep], -1)[0]
+        rest_selected = mx.split(selected_logprobs, [min_tokens_to_keep], -1)[1]
+        selected_logprobs = mx.concatenate([top_sorted, rest_selected], -1)
+      end
+
+      # Create inverse mapping to restore original order
+      inverse_indices = mx.put_along_axis(
+        mx.zeros_like(sorted_indices),
+        sorted_indices,
+        mx.arange(sorted_indices.shape[-1]).astype(sorted_indices.dtype),
+        -1
+      )
+
+      mx.take_along_axis(selected_logprobs, inverse_indices, -1)
+    end
+
+    def apply_top_p(logprobs, top_p)
+      mx = MLX::Core
+      probs = mx.exp(logprobs)
+      # sort in ascending order
+      sorted_indices = mx.argsort(logprobs, -1)
+      sorted_probs = mx.take_along_axis(probs, sorted_indices, -1)
+
+      cumulative_probs = mx.cumsum(sorted_probs, -1)
+
+      # Rearrange cumulative probs back to original order
+      inverse_indices = mx.put_along_axis(
+        mx.zeros_like(sorted_indices),
+        sorted_indices,
+        mx.arange(sorted_indices.shape[-1]).astype(sorted_indices.dtype),
+        -1
+      )
+      cumulative_probs = mx.take_along_axis(cumulative_probs, inverse_indices, -1)
+
+      # select tokens with cumulative probs above threshold
+      threshold = mx.array(1.0 - top_p, dtype: cumulative_probs.dtype)
+      mask = mx.greater(cumulative_probs, threshold)
+      neg_inf = mx.array(-Float::INFINITY, dtype: logprobs.dtype)
+      mx.where(mask, logprobs, neg_inf)
+    end
+
+    def categorical_sampling(logits, temp)
+      mx = MLX::Core
+      mx.categorical(logits * (1.0 / temp))
+    end
+
+    def make_repetition_penalty(penalty, context_size = 20)
+      mx = MLX::Core
+      raise ArgumentError, "penalty must be a non-negative float" unless penalty.is_a?(Numeric) && penalty >= 0
+
+      ->(tokens, logits) {
+        if tokens && tokens.size > 0
+          recent = if tokens.is_a?(::Array)
+            tokens.last(context_size)
+          elsif tokens.respond_to?(:tolist)
+            tokens.tolist.last(context_size)
+          else
+            []
+          end
+          if recent.length > 0
+            token_indices = mx.array(recent, dtype: mx.int32)
+            n_tokens = recent.length
+            idx_2d = token_indices.reshape([1, n_tokens])
+            selected_logits = mx.take_along_axis(logits, idx_2d, -1)
+            zero = mx.array(0.0, dtype: selected_logits.dtype)
+            is_negative = mx.less(selected_logits, zero)
+            selected_logits = mx.where(
+              is_negative,
+              selected_logits * penalty,
+              selected_logits / penalty
+            )
+            logits = mx.put_along_axis(logits, idx_2d, selected_logits, -1)
+          end
+        end
+        logits
+      }
+    end
+  end
+end

data/lib/mlx_lm/server.rb

@@ -0,0 +1,190 @@
+require "json"
+require "securerandom"
+
+module MlxLm
+  module Server
+    # Request schema for POST /v1/chat/completions
+    class ChatCompletionRequest
+      attr_reader :model, :messages, :max_tokens, :temperature, :top_p, :stream, :stop
+
+      def self.from_hash(h)
+        new(
+          model: h["model"],
+          messages: h["messages"] || [],
+          max_tokens: h["max_tokens"] || 256,
+          temperature: h["temperature"] || 0.0,
+          top_p: h["top_p"] || 1.0,
+          stream: h.fetch("stream", false),
+          stop: h["stop"]
+        )
+      end
+
+      def initialize(model:, messages:, max_tokens: 256, temperature: 0.0, top_p: 1.0, stream: false, stop: nil)
+        @model = model
+        @messages = messages
+        @max_tokens = max_tokens
+        @temperature = temperature
+        @top_p = top_p
+        @stream = stream
+        @stop = stop
+      end
+    end
+
+    # Response schema for non-streaming chat completion
+    class ChatCompletionResponse
+      def initialize(model:, content:, prompt_tokens:, completion_tokens:, finish_reason: "stop")
+        @model = model
+        @content = content
+        @prompt_tokens = prompt_tokens
+        @completion_tokens = completion_tokens
+        @finish_reason = finish_reason
+        @id = "chatcmpl-#{SecureRandom.hex(12)}"
+        @created = Time.now.to_i
+      end
+
+      def to_hash
+        {
+          "id" => @id,
+          "object" => "chat.completion",
+          "created" => @created,
+          "model" => @model,
+          "choices" => [
+            {
+              "index" => 0,
+              "message" => {
+                "role" => "assistant",
+                "content" => @content,
+              },
+              "finish_reason" => @finish_reason,
+            }
+          ],
+          "usage" => {
+            "prompt_tokens" => @prompt_tokens,
+            "completion_tokens" => @completion_tokens,
+            "total_tokens" => @prompt_tokens + @completion_tokens,
+          }
+        }
+      end
+
+      def to_json
+        JSON.generate(to_hash)
+      end
+    end
+
+    # Streaming chunk response
+    class ChatCompletionChunk
+      def initialize(model:, content:, finish_reason: nil)
+        @model = model
+        @content = content
+        @finish_reason = finish_reason
+        @id = "chatcmpl-#{SecureRandom.hex(12)}"
+        @created = Time.now.to_i
+      end
+
+      def to_hash
+        {
+          "id" => @id,
+          "object" => "chat.completion.chunk",
+          "created" => @created,
+          "model" => @model,
+          "choices" => [
+            {
+              "index" => 0,
+              "delta" => {
+                "content" => @content,
+              },
+              "finish_reason" => @finish_reason,
+            }
+          ]
+        }
+      end
+
+      def to_sse
+        "data: #{JSON.generate(to_hash)}\n\n"
+      end
+    end
+
+    # GET /v1/models response
+    class ModelsListResponse
+      def initialize(models:)
+        @models = models
+      end
+
+      def to_hash
+        {
+          "object" => "list",
+          "data" => @models.map { |m|
+            {
+              "id" => m,
+              "object" => "model",
+              "created" => Time.now.to_i,
+              "owned_by" => "mlx-lm",
+            }
+          }
+        }
+      end
+
+      def to_json
+        JSON.generate(to_hash)
+      end
+    end
+
+    module_function
+
+    def start(model_path:, host: "127.0.0.1", port: 8080)
+      require "webrick"
+
+      model, tokenizer = LoadUtils.load(model_path)
+
+      server = WEBrick::HTTPServer.new(Port: port, BindAddress: host)
+
+      server.mount_proc "/v1/models" do |req, res|
+        res["Content-Type"] = "application/json"
+        resp = ModelsListResponse.new(models: [model_path])
+        res.body = resp.to_json
+      end
+
+      server.mount_proc "/v1/chat/completions" do |req, res|
+        body = JSON.parse(req.body)
+        chat_req = ChatCompletionRequest.from_hash(body)
+
+        prompt = ChatTemplate.apply(chat_req.messages)
+        sampler = SampleUtils.make_sampler(temp: chat_req.temperature, top_p: chat_req.top_p)
+
+        if chat_req.stream
+          res["Content-Type"] = "text/event-stream"
+          res["Cache-Control"] = "no-cache"
+
+          res.body = Enumerator.new { |yielder|
+            Generate.stream_generate(model, tokenizer, prompt,
+              max_tokens: chat_req.max_tokens, sampler: sampler).each do |resp|
+              chunk = ChatCompletionChunk.new(
+                model: chat_req.model,
+                content: resp.text,
+                finish_reason: resp.finish_reason
+              )
+              yielder << chunk.to_sse
+            end
+            yielder << "data: [DONE]\n\n"
+          }
+        else
+          text = Generate.generate(model, tokenizer, prompt,
+            max_tokens: chat_req.max_tokens, sampler: sampler)
+
+          res["Content-Type"] = "application/json"
+          resp = ChatCompletionResponse.new(
+            model: chat_req.model,
+            content: text,
+            prompt_tokens: prompt.length,
+            completion_tokens: text.length,
+            finish_reason: "stop"
+          )
+          res.body = resp.to_json
+        end
+      end
+
+      trap("INT") { server.shutdown }
+      server.start
+    end
+  end
+end

data/lib/mlx_lm/tokenizer_utils.rb

@@ -0,0 +1,158 @@
+require "tokenizers"
+require "json"
+
+module MlxLm
+  # Wraps a HuggingFace tokenizer (loaded via the tokenizers gem)
+  # providing encode/decode and metadata access.
+  class TokenizerWrapper
+    attr_reader :tokenizer
+
+    # Can be initialized with:
+    # 1. A path string (directory containing tokenizer.json)
+    # 2. A Tokenizers::Tokenizer object (with optional eos_token/eos_token_id)
+    def initialize(path_or_tokenizer, eos_token: nil, eos_token_id: nil)
+      if path_or_tokenizer.is_a?(String)
+        tokenizer_json = File.join(path_or_tokenizer, "tokenizer.json")
+        @tokenizer = Tokenizers::Tokenizer.from_file(tokenizer_json)
+
+        config_path = File.join(path_or_tokenizer, "tokenizer_config.json")
+        @config = File.exist?(config_path) ? JSON.parse(File.read(config_path)) : {}
+      else
+        @tokenizer = path_or_tokenizer
+        @config = {}
+      end
+
+      @eos_token_override = eos_token
+      @eos_token_id_override = eos_token_id
+
+      @_detokenizer = nil
+    end
+
+    def encode(text, add_special_tokens: true)
+      @tokenizer.encode(text, add_special_tokens: add_special_tokens).ids
+    end
+
+    def decode(ids, skip_special_tokens: false)
+      @tokenizer.decode(ids, skip_special_tokens: skip_special_tokens)
+    end
+
+    def eos_token
+      return @eos_token_override if @eos_token_override
+      token = @config["eos_token"]
+      token = token["content"] if token.is_a?(Hash)
+      token
+    end
+
+    def eos_token_id
+      # Try override ids first
+      if @eos_token_id_override && !@eos_token_id_override.empty?
+        return @eos_token_id_override.first
+      end
+
+      # Try config
+      if @config["eos_token"]
+        token = @config["eos_token"]
+        token = token["content"] if token.is_a?(Hash)
+        id = @tokenizer.token_to_id(token)
+        return id if id
+      end
+
+      # Try eos_token string override
+      if @eos_token_override
+        id = @tokenizer.token_to_id(@eos_token_override)
+        return id if id
+      end
+
+      nil
+    end
+
+    # Returns a Set of all EOS token IDs
+    def eos_token_ids
+      ids = Set.new
+      if @eos_token_id_override
+        @eos_token_id_override.each { |id| ids << id if id }
+      end
+      base_id = eos_token_id
+      ids << base_id if base_id
+      ids
+    end
+
+    def bos_token
+      token = @config["bos_token"]
+      token = token["content"] if token.is_a?(Hash)
+      token
+    end
+
+    def bos_token_id
+      if @config["bos_token"]
+        token = @config["bos_token"]
+        token = token["content"] if token.is_a?(Hash)
+        id = @tokenizer.token_to_id(token)
+        return id if id
+      end
+      nil
+    end
+
+    def vocab_size
+      @tokenizer.vocab_size
+    end
+
+    def id_to_token(id)
+      @tokenizer.id_to_token(id)
+    end
+
+    def token_to_id(token)
+      @tokenizer.token_to_id(token)
+    end
+
+    def detokenizer
+      @_detokenizer ||= StreamingDetokenizer.new(self)
+    end
+
+    def has_chat_template
+      !!@config["chat_template"]
+    end
+  end
+
+  # Streaming detokenizer that incrementally decodes tokens without O(T^2) cost.
+  # Uses a simple approach: maintain a buffer of token IDs, decode the full buffer,
+  # and emit only the new characters since the last decode.
+  class StreamingDetokenizer
+    attr_reader :last_segment
+
+    def initialize(tokenizer_wrapper)
+      @tokenizer = tokenizer_wrapper
+      @token_ids = []
+      @prev_text = ""
+      @last_segment = ""
+    end
+
+    # Add a token and record the new text segment
+    def add_token(token_id)
+      @token_ids << token_id
+      current_text = @tokenizer.decode(@token_ids)
+      @last_segment = current_text[@prev_text.length..] || ""
+      @prev_text = current_text
+      @last_segment
+    end
+
+    # Finalize and record any remaining text
+    def finalize
+      return "" if @token_ids.empty?
+      final = @tokenizer.decode(@token_ids)
+      @last_segment = final[@prev_text.length..] || ""
+      @prev_text = final
+      @last_segment
+    end
+
+    def text
+      @prev_text
+    end
+
+    def reset
+      @token_ids = []
+      @prev_text = ""
+      @last_segment = ""
+    end
+  end
+end

data/lib/mlx_lm/tuner/lora.rb

@@ -0,0 +1,165 @@
+module MlxLm
+  module Tuner
+    # LoRA adapter for Linear layers.
+    # Forward: y = linear(x) + scale * (dropout(x) @ lora_a @ lora_b)
+    class LoRALinear < MLX::NN::Module
+      def self.from_base(linear, r: 8, dropout: 0.0, scale: 20.0)
+        if linear.is_a?(MLX::NN::QuantizedLinear)
+          input_dims = linear.instance_variable_get(:@weight).shape[1] * 32 /
+                       (linear.instance_variable_get(:@bits) || 4)
+          output_dims = linear.instance_variable_get(:@weight).shape[0]
+          bias = !linear.instance_variable_get(:@bias).nil?
+        else
+          weight = linear.weight
+          output_dims, input_dims = weight.shape
+          bias = !linear.respond_to?(:bias) || !linear.bias.nil? rescue false
+        end
+        lora = new(input_dims, output_dims, r: r, dropout: dropout, scale: scale, bias: bias)
+        lora.linear = linear
+        lora
+      end
+
+      def initialize(input_dims, output_dims, r: 8, dropout: 0.0, scale: 20.0, bias: false)
+        super()
+        mx = MLX::Core
+        @scale = scale
+        self.linear = MLX::NN::Linear.new(input_dims, output_dims, bias: bias)
+        self.dropout = MLX::NN::Dropout.new(dropout)
+
+        # Initialize LoRA matrices
+        lora_scale = 1.0 / Math.sqrt(input_dims)
+        self.lora_a = mx.random_uniform(
+          [input_dims, r], -lora_scale, lora_scale, mx.float32
+        )
+        self.lora_b = mx.zeros([r, output_dims])
+      end
+
+      def call(x)
+        mx = MLX::Core
+        y = linear.call(x)
+        z = dropout.call(x)
+        z = mx.matmul(mx.matmul(z, lora_a), lora_b)
+        y + z * @scale
+      end
+
+      def fuse(dequantize: false)
+        mx = MLX::Core
+        lin = linear
+
+        if dequantize && lin.is_a?(MLX::NN::QuantizedLinear)
+          lin = MlxLm::Quantize.linear_from_quantized(lin)
+        end
+
+        weight = lin.weight
+        bias_val = lin.respond_to?(:bias) ? lin.bias : nil
+
+        # Fuse: W' = W + scale * (lora_a @ lora_b)^T
+        lora_weight = mx.matmul(lora_a, lora_b)
+        fused_weight = weight + mx.transpose(lora_weight) * @scale
+
+        out_features, in_features = fused_weight.shape
+        result = MLX::NN::Linear.new(in_features, out_features, bias: !bias_val.nil?)
+        result.weight = fused_weight
+        result.bias = bias_val if bias_val
+        result
+      end
+    end
+
+    # LoRA adapter for Embedding layers.
+    class LoRAEmbedding < MLX::NN::Module
+      def self.from_base(embedding, r: 8, dropout: 0.0, scale: 20.0)
+        weight = embedding.weight
+        num_embeddings, dims = weight.shape
+        lora = new(num_embeddings, dims, r: r, dropout: dropout, scale: scale)
+        lora.embedding = embedding
+        lora
+      end
+
+      def initialize(num_embeddings, dims, r: 8, dropout: 0.0, scale: 20.0)
+        super()
+        mx = MLX::Core
+        @scale = scale
+        self.embedding = MLX::NN::Embedding.new(num_embeddings, dims)
+        self.dropout = MLX::NN::Dropout.new(dropout)
+
+        lora_scale = 1.0 / Math.sqrt(num_embeddings)
+        self.lora_a = mx.random_uniform(
+          [num_embeddings, r], -lora_scale, lora_scale, mx.float32
+        )
+        self.lora_b = mx.zeros([r, dims])
+      end
+
+      def call(x)
+        mx = MLX::Core
+        y = embedding.call(x)
+        # LoRA for embedding: look up lora_a rows, then multiply by lora_b
+        z = mx.matmul(mx.take(lora_a, x, 0), lora_b)
+        z = dropout.call(z)
+        y + z * @scale
+      end
+
+      def as_linear(x)
+        mx = MLX::Core
+        y = embedding.as_linear(x)
+        z = mx.matmul(mx.matmul(dropout.call(x), mx.transpose(lora_b)), mx.transpose(lora_a))
+        y + z * @scale
+      end
+
+      def fuse(dequantize: false)
+        mx = MLX::Core
+        embed = embedding
+
+        if dequantize && embed.is_a?(MLX::NN::QuantizedEmbedding)
+          embed = MlxLm::Quantize.embedding_from_quantized(embed)
+        end
+
+        weight = embed.weight
+        lora_weight = mx.matmul(lora_a, lora_b)
+        fused_weight = weight + lora_weight * @scale
+
+        num_embeddings, dims = fused_weight.shape
+        result = MLX::NN::Embedding.new(num_embeddings, dims)
+        result.weight = fused_weight
+        result
+      end
+    end
+
+    module_function
+
+    # Default LoRA target keys (layer names that get LoRA applied)
+    DEFAULT_LORA_KEYS = %w[self_attn.q_proj self_attn.k_proj self_attn.v_proj].freeze
+
+    # Apply LoRA layers to a model's last N layers.
+    def apply_lora_layers(model, num_layers: nil, config: {})
+      rank = config["rank"] || config[:rank] || 8
+      scale = config["scale"] || config[:scale] || 20.0
+      dropout = config["dropout"] || config[:dropout] || 0.0
+      keys = config["keys"] || config[:keys] || DEFAULT_LORA_KEYS
+
+      layers = model.layers
+      num_layers ||= layers.length
+      target_layers = layers.last(num_layers)
+
+      target_layers.each do |layer|
+        _apply_lora_to_module(layer, "", keys, rank: rank, scale: scale, dropout: dropout)
+      end
+    end
+
+    def _apply_lora_to_module(mod, prefix, keys, rank:, scale:, dropout:)
+      mod.state.each do |key, value|
+        full_key = prefix.empty? ? key : "#{prefix}.#{key}"
+
+        if value.is_a?(MLX::NN::Linear) && keys.any? { |k| full_key.end_with?(k) || full_key.include?(k) }
+          lora = LoRALinear.from_base(value, r: rank, scale: scale, dropout: dropout)
+          mod.state[key] = lora
+        elsif value.is_a?(MLX::NN::Embedding) && keys.any? { |k| full_key.end_with?(k) || full_key.include?(k) }
+          lora = LoRAEmbedding.from_base(value, r: rank, scale: scale, dropout: dropout)
+          mod.state[key] = lora
+        elsif value.is_a?(MLX::NN::Module)
+          _apply_lora_to_module(value, full_key, keys, rank: rank, scale: scale, dropout: dropout)
+        end
+      end
+    end
+    module_function :_apply_lora_to_module
+  end
+end

data/lib/mlx_lm/version.rb

@@ -0,0 +1,3 @@
+module MlxLm
+  VERSION = "0.30.7.1"
+end