mlx-ruby-lm 0.30.7.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 9a305e7fa3e70c61d29f8e997a98d4f7da6ff344d2df698778f56980dae27a34
+  data.tar.gz: f81553b3050391f1585f5ba6822c1cd28936f5ddaccc0b852bcfae30ea41bdf4
+SHA512:
+  metadata.gz: d5e6be331e95c5323b7fdb93b8defb2242a5ad969269e8e5442cc8f499d1fb2beabe03f81d14e7ee8abc3c61bf5d144ea5755bda7462cfdb9bcddb0afe9695f0
+  data.tar.gz: b333ba48c0e390b217ec4320e149f03b1e507caefe856701bf85245fa301b55d0fec2d75ecc2bdd92e367b49e3c82454b1b0baffe1d9c161223fb9580635a20e

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2025 Alex Skryl
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,83 @@
+# mlx-ruby-lm
+
+Ruby LLM inference toolkit built on the `mlx` gem.
+
+## Index
+
+- [Documentation Index](docs/index.md)
+- [Installation](docs/installation.md)
+- [CLI Usage](docs/cli.md)
+- [Ruby APIs](docs/ruby-apis.md)
+- [Models](docs/models.md)
+
+For full reference pages and deep dives, start at [docs/index.md](docs/index.md).
+
+## Installation
+
+```bash
+gem install mlx-ruby-lm
+```
+
+Or add it to a project:
+
+```bash
+bundle add mlx-ruby-lm
+```
+
+See [docs/installation.md](docs/installation.md) for requirements and source installs.
+
+## CLI Usage
+
+Executable: `mlx_lm`
+
+Commands:
+
+- `mlx_lm generate`
+- `mlx_lm chat`
+- `mlx_lm server`
+
+Quick examples:
+
+```bash
+mlx_lm generate --model /path/to/model --prompt "Hello"
+mlx_lm chat --model /path/to/model --system-prompt "You are concise."
+mlx_lm server --model /path/to/model --host 127.0.0.1 --port 8080
+```
+
+See [docs/cli.md](docs/cli.md) for options, defaults, and current parser/behavior caveats.
+
+## High-Level Ruby API Usage
+
+```ruby
+require "mlx"
+require "mlx_lm"
+
+model, tokenizer = MlxLm::LoadUtils.load("/path/to/model")
+text = MlxLm::Generate.generate(model, tokenizer, "Hello", max_tokens: 64)
+puts text
+```
+
+Streaming:
+
+```ruby
+MlxLm::Generate.stream_generate(model, tokenizer, "Hello", max_tokens: 64).each do |resp|
+  print resp.text
+end
+puts
+```
+
+See [docs/ruby-apis.md](docs/ruby-apis.md) for the full API inventory.
+
+## High-Level Model Usage
+
+`LoadUtils.load` expects a local model directory with files such as `config.json`,
+`tokenizer.json`, and `model*.safetensors`.
+
+To inspect supported model keys at runtime:
+
+```ruby
+require "mlx_lm"
+puts MlxLm::Models::REGISTRY.keys.sort
+```
+
+See [docs/models.md](docs/models.md) for full registry keys and remapping behavior.

data/exe/mlx_lm

@@ -0,0 +1,7 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "mlx"
+require "mlx_lm"
+
+MlxLm::CLI.run(ARGV)

data/lib/mlx_lm/benchmark.rb

@@ -0,0 +1,67 @@
+module MlxLm
+  module Benchmark
+    module_function
+
+    # Measure generation performance (tokens/sec).
+    def measure_generation(model, prompt_tokens: 32, gen_tokens: 64, vocab_size: 32000)
+      mx = MLX::Core
+
+      # Create random prompt tokens
+      prompt = mx.random_uniform([prompt_tokens], 0.0, (vocab_size - 1).to_f, mx.float32).astype(mx.int32)
+      mx.eval(prompt)
+
+      # Create cache
+      cache = Cache.make_prompt_cache(model)
+
+      # Measure prompt processing
+      prompt_input = prompt.reshape([1, prompt_tokens])
+      prompt_start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      logits = model.call(prompt_input, cache: cache)
+      mx.eval(logits)
+      mx.eval(*cache.map(&:state).flatten.compact)
+      prompt_elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - prompt_start
+      prompt_tps = prompt_tokens.to_f / [prompt_elapsed, 1e-9].max
+
+      # Get first generated token
+      last_logits = logits.reshape([prompt_tokens, logits.shape[-1]])
+      # Take last position
+      last_pos = mx.split(last_logits, [prompt_tokens - 1], 0)[1]
+      y = mx.argmax(last_pos, -1)
+      mx.eval(y)
+
+      # Measure generation
+      gen_start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      gen_tokens.times do
+        y_input = y.reshape([1, 1])
+        logits = model.call(y_input, cache: cache)
+        mx.eval(logits)
+        mx.eval(*cache.map(&:state).flatten.compact)
+        y = mx.argmax(logits.reshape([1, logits.shape[-1]]), -1)
+        mx.eval(y)
+      end
+      gen_elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - gen_start
+      gen_tps = gen_tokens.to_f / [gen_elapsed, 1e-9].max
+
+      {
+        prompt_tokens: prompt_tokens,
+        prompt_time: prompt_elapsed,
+        prompt_tps: prompt_tps,
+        generation_tokens: gen_tokens,
+        generation_time: gen_elapsed,
+        generation_tps: gen_tps,
+      }
+    end
+
+    # Get model statistics (parameter count, etc.)
+    def model_stats(model)
+      params = MLX::Utils.tree_flatten(model.parameters)
+      total = 0
+      params.each { |_, v| total += v.size }
+
+      {
+        total_params: total,
+        num_layers: model.respond_to?(:layers) ? model.layers.length : 0,
+      }
+    end
+  end
+end

data/lib/mlx_lm/chat_template.rb

@@ -0,0 +1,41 @@
+module MlxLm
+  module ChatTemplate
+    module_function
+
+    # Apply a simple chat template to format messages into a prompt string.
+    # This is a default/fallback template. Model-specific templates (like
+    # Jinja-based ones from tokenizer_config.json) can override this.
+    def apply(messages, template: :default)
+      case template
+      when :default
+        apply_default(messages)
+      when :chatml
+        apply_chatml(messages)
+      else
+        apply_default(messages)
+      end
+    end
+
+    # Default template: ChatML-like format
+    # <|im_start|>system
+    # content<|im_end|>
+    # <|im_start|>user
+    # content<|im_end|>
+    # <|im_start|>assistant
+    def apply_default(messages)
+      parts = []
+      messages.each do |msg|
+        role = msg["role"] || msg[:role]
+        content = msg["content"] || msg[:content]
+        parts << "<|im_start|>#{role}\n#{content}<|im_end|>"
+      end
+      parts << "<|im_start|>assistant"
+      parts.join("\n")
+    end
+
+    # ChatML template (same as default, widely used)
+    def apply_chatml(messages)
+      apply_default(messages)
+    end
+  end
+end

data/lib/mlx_lm/cli.rb

@@ -0,0 +1,113 @@
+require "optparse"
+
+module MlxLm
+  module CLI
+    COMMANDS = %w[generate chat server].freeze
+
+    module_function
+
+    def parse_args(argv)
+      command = argv.shift
+      unless COMMANDS.include?(command)
+        raise ArgumentError, "Unknown command '#{command}'. Valid commands: #{COMMANDS.join(', ')}"
+      end
+
+      args = default_args.merge(command: command)
+
+      parser = OptionParser.new do |opts|
+        opts.banner = "Usage: mlx_lm #{command} [options]"
+
+        opts.on("--model MODEL", "Model path or HuggingFace ID") { |v| args[:model] = v }
+        opts.on("--prompt PROMPT", "Input prompt") { |v| args[:prompt] = v }
+        opts.on("--max-tokens N", Integer, "Maximum tokens to generate") { |v| args[:max_tokens] = v }
+        opts.on("--temp TEMP", Float, "Sampling temperature") { |v| args[:temp] = v }
+        opts.on("--top-p P", Float, "Top-p (nucleus) sampling") { |v| args[:top_p] = v }
+        opts.on("--seed N", Integer, "Random seed") { |v| args[:seed] = v }
+        opts.on("--repetition-penalty F", Float, "Repetition penalty") { |v| args[:repetition_penalty] = v }
+        opts.on("--repetition-context-size N", Integer, "Repetition context size") { |v| args[:repetition_context_size] = v }
+        opts.on("--host HOST", "Server host") { |v| args[:host] = v }
+        opts.on("--port PORT", Integer, "Server port") { |v| args[:port] = v }
+        opts.on("--system-prompt PROMPT", "System prompt for chat") { |v| args[:system_prompt] = v }
+        opts.on("--verbose", "Verbose output") { args[:verbose] = true }
+      end
+
+      parser.parse!(argv)
+      args
+    end
+
+    def default_args
+      {
+        command: nil,
+        model: nil,
+        prompt: "",
+        max_tokens: 256,
+        temp: 0.0,
+        top_p: 1.0,
+        seed: nil,
+        repetition_penalty: nil,
+        repetition_context_size: 20,
+        host: "127.0.0.1",
+        port: 8080,
+        system_prompt: nil,
+        verbose: false,
+      }
+    end
+
+    def run(argv = ARGV)
+      args = parse_args(argv.dup)
+
+      case args[:command]
+      when "generate"
+        run_generate(args)
+      when "chat"
+        run_chat(args)
+      when "server"
+        run_server(args)
+      end
+    end
+
+    def run_generate(args)
+      model, tokenizer = LoadUtils.load(args[:model])
+      sampler = SampleUtils.make_sampler(temp: args[:temp], top_p: args[:top_p])
+      text = Generate.generate(model, tokenizer, args[:prompt],
+        max_tokens: args[:max_tokens], sampler: sampler, verbose: args[:verbose])
+      puts text unless args[:verbose]
+    end
+
+    def run_chat(args)
+      model, tokenizer = LoadUtils.load(args[:model])
+      messages = []
+      if args[:system_prompt]
+        messages << { "role" => "system", "content" => args[:system_prompt] }
+      end
+
+      loop do
+        print "> "
+        $stdout.flush
+        input = $stdin.gets
+        break if input.nil?
+        input = input.strip
+        break if input.empty? || input == "exit" || input == "quit"
+
+        messages << { "role" => "user", "content" => input }
+        prompt = ChatTemplate.apply(messages)
+
+        sampler = SampleUtils.make_sampler(temp: args[:temp])
+        text = ""
+        Generate.stream_generate(model, tokenizer, prompt,
+          max_tokens: args[:max_tokens], sampler: sampler).each do |resp|
+          print resp.text
+          $stdout.flush
+          text += resp.text
+        end
+        puts
+
+        messages << { "role" => "assistant", "content" => text }
+      end
+    end
+
+    def run_server(args)
+      Server.start(model_path: args[:model], host: args[:host], port: args[:port])
+    end
+  end
+end

data/lib/mlx_lm/config.rb

@@ -0,0 +1,30 @@
+require "json"
+
+module MlxLm
+  module Config
+    module_function
+
+    # Load model config from a directory containing config.json
+    # and optionally generation_config.json.
+    # Mirrors Python mlx_lm.utils.load_config
+    def load(model_path)
+      config_path = File.join(model_path, "config.json")
+      config = JSON.parse(File.read(config_path))
+
+      gen_config_path = File.join(model_path, "generation_config.json")
+      if File.exist?(gen_config_path)
+        begin
+          gen_config = JSON.parse(File.read(gen_config_path))
+        rescue JSON::ParserError
+          gen_config = {}
+        end
+
+        if (eos = gen_config["eos_token_id"])
+          config["eos_token_id"] = eos
+        end
+      end
+
+      config
+    end
+  end
+end

data/lib/mlx_lm/convert_utils.rb

@@ -0,0 +1,51 @@
+module MlxLm
+  module ConvertUtils
+    DTYPE_MAP = {
+      "float32" => :float32,
+      "float16" => :float16,
+      "bfloat16" => :bfloat16,
+      "int8" => :int8,
+      "int32" => :int32,
+    }.freeze
+
+    module_function
+
+    # Convert an MLX array to a different dtype.
+    def convert_dtype(array, target_dtype)
+      if target_dtype.is_a?(MLX::Core::Dtype)
+        return array.astype(target_dtype)
+      end
+      dtype_sym = target_dtype.is_a?(Symbol) ? target_dtype : DTYPE_MAP[target_dtype.to_s]
+      raise ArgumentError, "Unknown dtype: #{target_dtype}" unless dtype_sym
+      array.astype(MLX::Core::Dtype.new(dtype_sym))
+    end
+
+    # Count total number of parameters in a model.
+    def count_parameters(model)
+      params = MLX::Utils.tree_flatten(model.parameters)
+      total = 0
+      params.each { |_, v| total += v.size }
+      total
+    end
+
+    # Estimate total model size in bytes.
+    def model_size_bytes(model)
+      mx = MLX::Core
+      params = MLX::Utils.tree_flatten(model.parameters)
+      total = 0
+      params.each do |_, v|
+        bytes_per_elem = case v.dtype
+                         when mx.float32 then 4
+                         when mx.float16, mx.bfloat16 then 2
+                         when mx.int32 then 4
+                         when mx.int8, mx.uint8 then 1
+                         when mx.int16, mx.uint16 then 2
+                         when mx.int64 then 8
+                         else 4
+                         end
+        total += v.size * bytes_per_elem
+      end
+      total
+    end
+  end
+end

data/lib/mlx_lm/generate.rb

@@ -0,0 +1,204 @@
+module MlxLm
+  # Response object yielded during streaming generation
+  GenerationResponse = Struct.new(
+    :text,
+    :token,
+    :logprobs,
+    :prompt_tokens,
+    :prompt_tps,
+    :generation_tokens,
+    :generation_tps,
+    :peak_memory,
+    :finish_reason,
+    keyword_init: true
+  )
+
+  module Generate
+    module_function
+
+    # A generator producing token ids based on the given prompt from the model.
+    # Yields [token_id, logprobs] for each generated token.
+    def generate_step(
+      prompt,
+      model,
+      max_tokens: 256,
+      sampler: nil,
+      logits_processors: nil,
+      max_kv_size: nil,
+      prompt_cache: nil,
+      prefill_step_size: 2048
+    )
+      mx = MLX::Core
+
+      raise ArgumentError, "prompt must not be empty" if prompt.size == 0
+
+      tokens = nil
+
+      # Create the KV cache for generation
+      prompt_cache ||= Cache.make_prompt_cache(model, max_kv_size: max_kv_size)
+
+      sampler ||= ->(x) { mx.argmax(x, -1) }
+
+      model_call = ->(input_tokens_2d) {
+        model.call(input_tokens_2d, cache: prompt_cache)
+      }
+
+      step = ->(input_tokens_1d) {
+        seq_len = input_tokens_1d.size
+        input_2d = input_tokens_1d.reshape([1, seq_len])
+        logits = model_call.call(input_2d)
+
+        # Take the last token's logits
+        last_dim = logits.shape[1]
+        if last_dim > 1
+          logits = mx.split(logits, [last_dim - 1], 1)[1]
+        end
+        vocab_size = logits.shape[-1]
+        logits = logits.reshape([1, vocab_size])
+
+        if logits_processors && input_tokens_1d.size > 0
+          tokens = if tokens.nil?
+            input_tokens_1d
+          else
+            mx.concatenate([tokens, input_tokens_1d], 0)
+          end
+          logits_processors.each { |processor| logits = processor.call(tokens, logits) }
+        end
+
+        logprobs = logits - mx.logsumexp(logits, -1, true)
+        sampled = sampler.call(logprobs)
+        [sampled, logprobs.reshape([vocab_size])]
+      }
+
+      # Prompt prefilling - process prompt in chunks
+      prompt_arr = prompt.is_a?(::Array) ? mx.array(prompt, dtype: mx.uint32) : prompt
+      total_prompt_tokens = prompt_arr.size
+
+      # Process prompt chunks (all but last token)
+      while total_prompt_tokens > 1
+        remaining = total_prompt_tokens - 1
+        n_to_process = [prefill_step_size, remaining].min
+        chunk = mx.split(prompt_arr, [n_to_process], 0)[0]
+        chunk_len = chunk.size
+        model_call.call(chunk.reshape([1, chunk_len]))
+        mx.eval(*prompt_cache.map(&:state).flatten.compact)
+        prompt_arr = mx.split(prompt_arr, [n_to_process], 0)[1]
+        total_prompt_tokens -= n_to_process
+      end
+
+      # Process last token and get first generated token
+      y, logprobs = step.call(prompt_arr)
+      mx.eval(y, logprobs)
+
+      Enumerator.new do |yielder|
+        n = 0
+        loop do
+          break if n == max_tokens
+
+          y_1d = y.ndim > 1 ? y.reshape([y.size]) : y
+          next_y, next_logprobs = step.call(y_1d)
+          mx.eval(next_y, next_logprobs)
+
+          yielder.yield [y.item, logprobs]
+          y, logprobs = next_y, next_logprobs
+          n += 1
+        end
+      end
+    end
+
+    # Stream text generation from the model.
+    # Yields GenerationResponse objects with text segments.
+    def stream_generate(model, tokenizer, prompt, max_tokens: 256, **kwargs)
+      tokenizer = TokenizerWrapper.new(tokenizer) unless tokenizer.is_a?(TokenizerWrapper)
+
+      unless prompt.is_a?(MLX::Core::Array)
+        if prompt.is_a?(String)
+          prompt = tokenizer.encode(prompt)
+        end
+        prompt = MLX::Core.array(prompt, dtype: MLX::Core.uint32)
+      end
+
+      detokenizer = tokenizer.detokenizer
+
+      token_generator = generate_step(prompt, model, max_tokens: max_tokens, **kwargs)
+
+      tic = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      prompt_tps = 0.0
+
+      Enumerator.new do |yielder|
+        n = 0
+        last_token = nil
+        token_generator.each do |token, logprobs|
+          if n == 0
+            prompt_time = Process.clock_gettime(Process::CLOCK_MONOTONIC) - tic
+            prompt_tps = prompt.size.to_f / [prompt_time, 1e-9].max
+            tic = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          end
+
+          last_token = token
+
+          if tokenizer.eos_token_ids.include?(token)
+            detokenizer.finalize
+            elapsed = [Process.clock_gettime(Process::CLOCK_MONOTONIC) - tic, 1e-9].max
+            yielder.yield GenerationResponse.new(
+              text: detokenizer.last_segment,
+              token: token,
+              logprobs: logprobs,
+              prompt_tokens: prompt.size,
+              prompt_tps: prompt_tps,
+              generation_tokens: n + 1,
+              generation_tps: (n + 1).to_f / elapsed,
+              peak_memory: 0.0,
+              finish_reason: "stop"
+            )
+            break
+          end
+
+          detokenizer.add_token(token)
+          elapsed = [Process.clock_gettime(Process::CLOCK_MONOTONIC) - tic, 1e-9].max
+
+          yielder.yield GenerationResponse.new(
+            text: detokenizer.last_segment,
+            token: token,
+            logprobs: logprobs,
+            prompt_tokens: prompt.size,
+            prompt_tps: prompt_tps,
+            generation_tokens: n + 1,
+            generation_tps: (n + 1).to_f / elapsed,
+            peak_memory: 0.0,
+            finish_reason: ((n + 1) == max_tokens ? "length" : nil)
+          )
+
+          n += 1
+          break if (n + 1) == max_tokens
+        end
+      end
+    end
+
+    # Non-streaming generation, returns complete text.
+    def generate(model, tokenizer, prompt, verbose: false, **kwargs)
+      text = ""
+      response = nil
+      stream_generate(model, tokenizer, prompt, **kwargs).each do |resp|
+        text += resp.text
+        response = resp
+        if verbose
+          print resp.text
+          $stdout.flush
+        end
+      end
+
+      if verbose
+        puts
+        puts "=" * 10
+        if text.empty?
+          puts "No text generated for this prompt"
+          return text
+        end
+        puts "Prompt: #{response.prompt_tokens} tokens, #{'%.3f' % response.prompt_tps} tokens-per-sec"
+        puts "Generation: #{response.generation_tokens} tokens, #{'%.3f' % response.generation_tps} tokens-per-sec"
+      end
+      text
+    end
+  end
+end