mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/telechat3.rb

@@ -0,0 +1,192 @@
+module MlxLm
+  module Models
+    module Telechat3
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "telechat3"
+        field :hidden_size, default: 4096
+        field :intermediate_size, default: 14336
+        field :max_position_embeddings, default: 32768
+        field :num_attention_heads, default: 32
+        field :num_hidden_layers, default: 32
+        field :num_key_value_heads, default: nil
+        field :rms_norm_eps, default: 1e-6
+        field :vocab_size, default: 151936
+        field :rope_theta, default: 10_000.0
+        field :mlp_bias, default: false
+        field :attention_bias, default: false
+        field :head_dim, default: nil
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: false
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+        end
+      end
+
+      class Telechat3Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          @num_attention_heads = args.num_attention_heads
+          @num_key_value_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(
+            dim,
+            args.num_attention_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.k_proj = MLX::NN::Linear.new(
+            dim,
+            args.num_key_value_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.v_proj = MLX::NN::Linear.new(
+            dim,
+            args.num_key_value_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.o_proj = MLX::NN::Linear.new(
+            args.num_attention_heads * @head_dim,
+            dim,
+            bias: args.attention_bias
+          )
+
+          self.rope = MlxLm::Models.initialize_rope(
+            @head_dim,
+            args.rope_theta,
+            false,
+            args.rope_scaling,
+            max_position_embeddings: args.max_position_embeddings
+          )
+        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, @num_attention_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([b, l, @num_key_value_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @num_key_value_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @num_attention_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class Telechat3MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            args.intermediate_size,
+            bias: args.mlp_bias
+          )
+          self.down_proj = MLX::NN::Linear.new(
+            args.intermediate_size,
+            args.hidden_size,
+            bias: args.mlp_bias
+          )
+          self.up_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            args.intermediate_size,
+            bias: args.mlp_bias
+          )
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(gate_proj.call(x), up_proj.call(x)))
+        end
+      end
+
+      class Telechat3DecoderLayer < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Telechat3Attention.new(args)
+          self.mlp = Telechat3MLP.new(args)
+          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
+          h + mlp.call(post_attention_layernorm.call(h))
+        end
+      end
+
+      class Telechat3Model < 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) { Telechat3DecoderLayer.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          h = input_embeddings || embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.model = Telechat3Model.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          out = model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          result = weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("telechat3", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/youtu_llm.rb

@@ -0,0 +1,230 @@
+module MlxLm
+  module Models
+    module YoutuLLM
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "youtu_llm"
+        field :vocab_size, default: 128_256
+        field :hidden_size, default: 2048
+        field :intermediate_size, default: 6144
+        field :num_hidden_layers, default: 32
+        field :num_attention_heads, default: 16
+        field :num_key_value_heads, default: 16
+        field :kv_lora_rank, default: 512
+        field :q_lora_rank, default: 1536
+        field :qk_rope_head_dim, default: 64
+        field :v_head_dim, default: 128
+        field :qk_nope_head_dim, default: 128
+        field :max_position_embeddings, default: 131_072
+        field :rms_norm_eps, default: 1e-6
+        field :rope_theta, default: 1_600_000.0
+        field :rope_traditional, default: true
+        field :rope_scaling, default: nil
+        field :attention_bias, default: false
+        field :mlp_bias, default: false
+        field :tie_word_embeddings, default: true
+      end
+
+      class YoutuLLMAttention < MLX::NN::Module
+        def initialize(config)
+          super()
+          @hidden_size = config.hidden_size
+          @num_heads = config.num_attention_heads
+          @q_lora_rank = config.q_lora_rank
+          @qk_rope_head_dim = config.qk_rope_head_dim
+          @kv_lora_rank = config.kv_lora_rank
+          @v_head_dim = config.v_head_dim
+          @qk_nope_head_dim = config.qk_nope_head_dim
+          @q_head_dim = @qk_nope_head_dim + @qk_rope_head_dim
+          @kv_head_dim = @qk_nope_head_dim + @v_head_dim
+          @scale = @q_head_dim**(-0.5)
+
+          if @q_lora_rank.nil?
+            self.q_proj = MLX::NN::Linear.new(
+              @hidden_size,
+              @num_heads * @q_head_dim,
+              bias: false
+            )
+          else
+            self.q_a_proj = MLX::NN::Linear.new(
+              @hidden_size,
+              @q_lora_rank,
+              bias: config.attention_bias
+            )
+            self.q_a_layernorm = MLX::NN::RMSNorm.new(@q_lora_rank, eps: config.rms_norm_eps)
+            self.q_b_proj = MLX::NN::Linear.new(@q_lora_rank, @num_heads * @q_head_dim, bias: false)
+          end
+
+          self.kv_a_proj_with_mqa = MLX::NN::Linear.new(
+            @hidden_size,
+            @kv_lora_rank + @qk_rope_head_dim,
+            bias: config.attention_bias
+          )
+          self.kv_a_layernorm = MLX::NN::RMSNorm.new(@kv_lora_rank, eps: config.rms_norm_eps)
+          self.kv_b_proj = MLX::NN::Linear.new(
+            @kv_lora_rank,
+            @num_heads * (@q_head_dim - @qk_rope_head_dim + @v_head_dim),
+            bias: false
+          )
+
+          self.o_proj = MLX::NN::Linear.new(
+            @num_heads * @v_head_dim,
+            @hidden_size,
+            bias: config.attention_bias
+          )
+
+          self.rope = MlxLm::Models.initialize_rope(
+            @qk_rope_head_dim,
+            config.rope_theta,
+            config.rope_traditional,
+            config.rope_scaling,
+            max_position_embeddings: config.max_position_embeddings
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          q = if @q_lora_rank.nil?
+            q_proj.call(x)
+          else
+            q_b_proj.call(q_a_layernorm.call(q_a_proj.call(x)))
+          end
+
+          q = q.reshape([b, l, @num_heads, @q_head_dim]).transpose([0, 2, 1, 3])
+          q_nope, q_pe = mx.split(q, [@qk_nope_head_dim], -1)
+
+          compressed_kv = kv_a_proj_with_mqa.call(x)
+          compressed_kv, k_pe = mx.split(compressed_kv, [@kv_lora_rank], -1)
+          k_pe = k_pe.reshape([b, l, 1, @qk_rope_head_dim]).transpose([0, 2, 1, 3])
+
+          kv = kv_b_proj.call(kv_a_layernorm.call(compressed_kv))
+          kv = kv.reshape([b, l, @num_heads, @kv_head_dim]).transpose([0, 2, 1, 3])
+          k_nope, values = mx.split(kv, [@qk_nope_head_dim], -1)
+
+          if cache
+            q_pe = rope.call(q_pe, offset: cache.offset)
+            k_pe = rope.call(k_pe, offset: cache.offset)
+            k_pe = mx.repeat(k_pe, @num_heads, 1)
+            keys, values = cache.update_and_fetch(mx.concatenate([k_nope, k_pe], -1), values)
+          else
+            q_pe = rope.call(q_pe)
+            k_pe = rope.call(k_pe)
+            k_pe = mx.repeat(k_pe, @num_heads, 1)
+            keys = mx.concatenate([k_nope, k_pe], -1)
+          end
+
+          queries = mx.concatenate([q_nope, q_pe], -1)
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @num_heads * @v_head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class YoutuLLMMLP < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(
+            config.hidden_size,
+            config.intermediate_size,
+            bias: config.mlp_bias
+          )
+          self.up_proj = MLX::NN::Linear.new(
+            config.hidden_size,
+            config.intermediate_size,
+            bias: config.mlp_bias
+          )
+          self.down_proj = MLX::NN::Linear.new(
+            config.intermediate_size,
+            config.hidden_size,
+            bias: config.mlp_bias
+          )
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(gate_proj.call(x), up_proj.call(x)))
+        end
+      end
+
+      class YoutuLLMDecoderLayer < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.self_attn = YoutuLLMAttention.new(config)
+          self.mlp = YoutuLLMMLP.new(config)
+          self.input_layernorm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.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 = mlp.call(post_attention_layernorm.call(h))
+          h + r
+        end
+      end
+
+      class YoutuLLMModel < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.embed_tokens = MLX::NN::Embedding.new(config.vocab_size, config.hidden_size)
+          self.layers = Array.new(config.num_hidden_layers) { YoutuLLMDecoderLayer.new(config) }
+          self.norm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.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(config)
+          super()
+          @config = config
+          self.model_type = config.model_type
+          self.model = YoutuLLMModel.new(config)
+          unless config.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(config.hidden_size, config.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @config.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          result = weights.dup
+          result.delete("lm_head.weight") if @config.tie_word_embeddings
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("youtu_llm", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models.rb

@@ -0,0 +1,33 @@
+module MlxLm
+  module Models
+    # Model registry: maps architecture name to [Model, ModelArgs] classes.
+    # Additional architectures register themselves here.
+    REGISTRY = {}
+
+    # Remapping for architectures that share implementation
+    REMAPPING = {
+      "mistral" => "llama",
+      "falcon_mamba" => "mamba",
+    }.freeze
+
+    module_function
+
+    def register(name, model_class, args_class)
+      REGISTRY[name] = [model_class, args_class]
+    end
+
+    def get_classes(config)
+      model_type = config["model_type"]
+      raise ArgumentError, "config.json missing 'model_type' field" unless model_type
+
+      # Apply remapping
+      canonical = REMAPPING.fetch(model_type, model_type)
+
+      unless REGISTRY.key?(canonical)
+        raise ArgumentError, "Model architecture '#{model_type}' (canonical: '#{canonical}') not found in registry. Available: #{REGISTRY.keys.join(', ')}"
+      end
+
+      REGISTRY[canonical]
+    end
+  end
+end

data/lib/mlx_lm/perplexity.rb

@@ -0,0 +1,48 @@
+module MlxLm
+  module Perplexity
+    module_function
+
+    # Compute perplexity of a model on a token sequence.
+    # Returns exp(average negative log-likelihood per token).
+    def compute(model, tokens, batch_size: nil)
+      ll = log_likelihood(model, tokens, batch_size: batch_size)
+      num_tokens = tokens.size - 1
+      avg_nll = -ll / num_tokens.to_f
+      Math.exp(avg_nll)
+    end
+
+    # Compute total log-likelihood of a token sequence.
+    # Sum of log P(token_i | token_0..token_{i-1}) for i in 1..n.
+    def log_likelihood(model, tokens, batch_size: nil)
+      mx = MLX::Core
+
+      token_arr = tokens.is_a?(MLX::Core::Array) ? tokens : mx.array(tokens, dtype: mx.int32)
+      total_tokens = token_arr.size
+
+      # Process all at once for small sequences
+      input = token_arr.reshape([1, total_tokens])
+      logits = model.call(input)
+      mx.eval(logits)
+
+      # Compute log probabilities
+      # logits shape: [1, total_tokens, vocab_size]
+      # We want P(token[i+1] | token[0..i])
+      vocab_size = logits.shape[-1]
+      logits_2d = logits.reshape([total_tokens, vocab_size])
+
+      # Log softmax
+      log_probs = logits_2d - mx.logsumexp(logits_2d, -1, true)
+
+      # Gather log probs for actual next tokens
+      # For position i, the model predicts token i+1
+      total_ll = 0.0
+      (0...(total_tokens - 1)).each do |i|
+        target_token = token_arr[i + 1].item
+        lp = log_probs[i][target_token].item
+        total_ll += lp
+      end
+
+      total_ll
+    end
+  end
+end

data/lib/mlx_lm/quantize.rb

@@ -0,0 +1,131 @@
+module MlxLm
+  module Quantize
+    module_function
+
+    # Quantize a model's linear and embedding layers.
+    #
+    # @param model [nn::Module] The model to quantize
+    # @param group_size [Integer] Group size for quantization (default: 64)
+    # @param bits [Integer] Number of bits (default: 4)
+    # @param weights [Hash] Optional weight dict; if provided, only quantize layers
+    #   that have corresponding .scales keys
+    # @return [Hash] The quantization config
+    def quantize_model(model, group_size: 64, bits: 4, weights: nil)
+      if weights
+        # Auto-detect: only quantize layers that have .scales in weights
+        class_predicate = ->(path, mod) {
+          if mod.respond_to?(:to_quantized)
+            weights.key?("#{path}.scales")
+          else
+            false
+          end
+        }
+      else
+        # Quantize all quantizable layers
+        class_predicate = ->(_path, mod) {
+          mod.respond_to?(:to_quantized)
+        }
+      end
+
+      MLX::NN.quantize(model, group_size: group_size, bits: bits, class_predicate: class_predicate)
+
+      { "group_size" => group_size, "bits" => bits }
+    end
+
+    # Dequantize a model (convert QuantizedLinear back to Linear, etc.)
+    #
+    # @param model [nn::Module] The quantized model to dequantize
+    # @return [nn::Module] The dequantized model
+    def dequantize_model(model)
+      mx = MLX::Core
+      de_quantize_layers(model)
+      model
+    end
+
+    # Compute bits per weight for a model
+    def bits_per_weight(model)
+      total_bits = 0
+      total_params = 0
+
+      model.named_modules.each do |name, mod|
+        case mod
+        when MLX::NN::QuantizedLinear
+          # Quantized: bits per element = quantized bits
+          weight = mod.instance_variable_get(:@weight)
+          if weight
+            num_params = weight.shape.reduce(:*)
+            total_bits += num_params * 4 # approximate
+            total_params += num_params
+          end
+        when MLX::NN::Linear
+          weight = mod.instance_variable_get(:@weight)
+          if weight
+            num_params = weight.shape.reduce(:*)
+            total_bits += num_params * 32 # float32
+            total_params += num_params
+          end
+        end
+      end
+
+      total_params > 0 ? total_bits.to_f / total_params : 0.0
+    end
+
+    private
+
+    def self.de_quantize_layers(mod)
+      mod.instance_variables.each do |ivar|
+        val = mod.instance_variable_get(ivar)
+        case val
+        when MLX::NN::QuantizedLinear
+          # Convert back to Linear
+          dequantized = linear_from_quantized(val)
+          mod.instance_variable_set(ivar, dequantized)
+        when MLX::NN::QuantizedEmbedding
+          dequantized = embedding_from_quantized(val)
+          mod.instance_variable_set(ivar, dequantized)
+        when MLX::NN::Module
+          de_quantize_layers(val)
+        when ::Array
+          val.each { |item| de_quantize_layers(item) if item.is_a?(MLX::NN::Module) }
+        end
+      end
+    end
+
+    def self.linear_from_quantized(qlinear)
+      mx = MLX::Core
+      weight = qlinear.instance_variable_get(:@weight)
+      scales = qlinear.instance_variable_get(:@scales)
+      biases = qlinear.instance_variable_get(:@biases)
+      bias = qlinear.instance_variable_get(:@bias)
+      group_size = qlinear.instance_variable_get(:@group_size) || 64
+      bits = qlinear.instance_variable_get(:@bits) || 4
+
+      # Dequantize weight
+      dequantized = mx.dequantize(weight, scales, biases, group_size, bits)
+      out_features = dequantized.shape[0]
+      in_features = dequantized.shape[1]
+
+      linear = MLX::NN::Linear.new(in_features, out_features, bias: !bias.nil?)
+      linear.instance_variable_set(:@weight, dequantized)
+      linear.instance_variable_set(:@bias, bias) if bias
+      linear
+    end
+
+    def self.embedding_from_quantized(qembed)
+      mx = MLX::Core
+      weight = qembed.instance_variable_get(:@weight)
+      scales = qembed.instance_variable_get(:@scales)
+      biases = qembed.instance_variable_get(:@biases)
+      group_size = qembed.instance_variable_get(:@group_size) || 64
+      bits = qembed.instance_variable_get(:@bits) || 4
+
+      dequantized = mx.dequantize(weight, scales, biases, group_size, bits)
+      num_embeddings = dequantized.shape[0]
+      dims = dequantized.shape[1]
+
+      embedding = MLX::NN::Embedding.new(num_embeddings, dims)
+      embedding.instance_variable_set(:@weight, dequantized)
+      embedding
+    end
+  end
+end