trainers-rb 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 3fd96717156bbd2f4d515f3b1a5aaea897f581b1888bb527124de305a6c2d536
+  data.tar.gz: 4a4a764def0528487ee20adb5e1edce52d7a3d267f1c99e399f5f48a7a5d9e2b
+SHA512:
+  metadata.gz: 4a36d98cfc6a61d3fe7cad411a653e242b7f2a11f402e7add4ea2300bd9ee59601c8408b864a1094ba6486444b4b9197c524d89de9b50e8c74ee263b3d59ab0b
+  data.tar.gz: da2ec1f4fd3d69e76f3d6a812aec49454ff9d99b99006e1b53f0086e4a60d1958d69b6a1b7447e033fcbf927c412cc3fea50954d0c3cd72e5ce49adf4140d95f

data/CHANGELOG.md

@@ -0,0 +1,17 @@
+# Changelog
+
+## 0.1.0 (2026-06-07)
+
+- Initial release
+- Trainer class with train/evaluate/predict
+- TrainingArguments configuration
+- AdamW optimizer with weight decay param groups
+- Linear, cosine, and constant learning rate schedulers with warmup
+- LoRA (Low-Rank Adaptation) for parameter-efficient fine-tuning
+- LoRA adapter save/load via safetensors
+- LoRA merge for inference deployment
+- DataCollatorWithPadding for dynamic batch padding
+- Callback system with PrinterCallback and EarlyStoppingCallback
+- Model saving via safetensors
+- CPU and MPS (Apple Silicon) device support
+- Gradient accumulation and gradient clipping

data/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Ruby ML Community
+
+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,293 @@
+# trainers-rb
+
+Fine-tune transformer models in Ruby.
+
+trainers-rb provides a training loop, LoRA (Low-Rank Adaptation), learning rate scheduling, and model serialization for HuggingFace transformer models loaded via [transformers-rb](https://github.com/ankane/transformers-ruby). It builds on [torch-rb](https://github.com/ankane/torch.rb) for autograd, optimizers, and tensor operations.
+
+All the heavy lifting happens in LibTorch C++ kernels. Ruby is the conductor.
+
+## Installation
+
+Add to your Gemfile:
+
+```ruby
+gem "trainers-rb"
+```
+
+Or install directly:
+
+```bash
+gem install trainers-rb
+```
+
+### Prerequisites
+
+trainers-rb depends on torch-rb, which requires LibTorch:
+
+```bash
+# macOS arm64
+curl -L -o /tmp/libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.4.0.zip
+unzip /tmp/libtorch.zip -d ~/libtorch
+bundle config set build.torch-rb --with-torch-dir=$HOME/libtorch/libtorch
+```
+
+## Quick Start
+
+```ruby
+require "trainers-rb"
+
+# Load a pre-trained model and tokenizer
+model, tokenizer = Trainers.from_pretrained(
+  "distilbert-base-uncased",
+  task: :sequence_classification,
+  num_labels: 2
+)
+
+# Prepare your dataset
+train_data = texts.map.with_index do |text, i|
+  encoded = tokenizer.(text, truncation: true, max_length: 128)
+  {
+    input_ids:      encoded["input_ids"],
+    attention_mask: encoded["attention_mask"],
+    labels:         labels[i]
+  }
+end
+train_dataset = Trainers::Dataset.new(train_data)
+
+# Configure and train
+args = Trainers::TrainingArguments.new(
+  output_dir:        "./output",
+  num_train_epochs:  3,
+  learning_rate:     2e-5,
+  eval_strategy:     :epoch
+)
+
+trainer = Trainers::Trainer.new(
+  model:         model,
+  args:          args,
+  train_dataset: train_dataset,
+  eval_dataset:  val_dataset,
+  tokenizer:     tokenizer,
+  data_collator: Trainers::DataCollatorWithPadding.new(tokenizer: tokenizer),
+  compute_metrics: ->(eval_pred) {
+    preds   = eval_pred.predictions.argmax(1)
+    correct = preds.eq(eval_pred.label_ids).sum.item
+    { accuracy: correct.to_f / eval_pred.label_ids.size(0) }
+  }
+)
+
+trainer.train
+trainer.save_model("./my-model")
+```
+
+## LoRA (Parameter-Efficient Fine-Tuning)
+
+Freeze 99% of parameters and train only small low-rank adapter matrices:
+
+```ruby
+# Apply LoRA to specific layers
+config = Trainers::LoraConfig.new(
+  r:              8,         # rank
+  lora_alpha:     16,        # scaling factor
+  lora_dropout:   0.1,
+  target_modules: ["query", "value"],  # which Linear layers to adapt
+  bias:           :none      # :none, :all, or :lora_only
+)
+
+Trainers::LoraModel.apply(model, config)
+# => LoRA applied to 12 modules: ...
+# => trainable params: 294,912 || all params: 66,955,010 || trainable%: 0.4404%
+
+# Train as usual
+trainer.train
+
+# Save just the adapters (tiny files)
+Trainers::LoraModel.save_adapters(model, "./lora-adapters")
+
+# Or merge into base model for inference
+Trainers::LoraModel.merge(model)
+trainer.save_model("./merged-model")
+```
+
+### Loading saved LoRA adapters
+
+```ruby
+model, tokenizer = Trainers.from_pretrained("distilbert-base-uncased", num_labels: 2)
+Trainers::LoraModel.apply(model, config)
+Trainers::LoraModel.load_adapters(model, "./lora-adapters")
+```
+
+## Training Arguments
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `output_dir` | `"./output"` | Directory for checkpoints and saved models |
+| `num_train_epochs` | `3` | Number of training epochs |
+| `per_device_train_batch_size` | `8` | Training batch size |
+| `per_device_eval_batch_size` | `8` | Evaluation batch size |
+| `learning_rate` | `5e-5` | Peak learning rate for AdamW |
+| `weight_decay` | `0.0` | Weight decay (applied to non-bias, non-norm params) |
+| `max_grad_norm` | `1.0` | Max gradient norm for clipping |
+| `gradient_accumulation_steps` | `1` | Accumulate gradients over N steps |
+| `warmup_steps` | `0` | Linear warmup steps |
+| `warmup_ratio` | `0.0` | Warmup as fraction of total steps (alternative to warmup_steps) |
+| `lr_scheduler_type` | `:linear` | `:linear`, `:cosine`, or `:constant` |
+| `eval_strategy` | `:no` | When to evaluate: `:no`, `:epoch`, or `:steps` |
+| `eval_steps` | `nil` | Evaluate every N steps (when `eval_strategy: :steps`) |
+| `save_strategy` | `:epoch` | When to save: `:no`, `:epoch`, or `:steps` |
+| `save_total_limit` | `nil` | Keep only the last N checkpoints |
+| `logging_steps` | `500` | Log every N steps |
+| `seed` | `42` | Random seed |
+| `no_mps` | `false` | Force CPU even if MPS is available |
+
+## Callbacks
+
+Built-in callbacks:
+
+```ruby
+# Early stopping
+early_stop = Trainers::EarlyStoppingCallback.new(
+  patience:    3,
+  threshold:   0.01,
+  metric_name: "eval_loss"
+)
+
+trainer = Trainers::Trainer.new(
+  model: model,
+  args: args,
+  callbacks: [early_stop],
+  # ...
+)
+```
+
+Custom callbacks:
+
+```ruby
+class WandbCallback < Trainers::TrainerCallback
+  def on_log(args, state, control, logs: nil, **)
+    # send logs to Weights & Biases, MLflow, etc.
+  end
+
+  def on_evaluate(args, state, control, metrics: nil, **)
+    # log evaluation metrics
+  end
+end
+```
+
+### Callback hooks
+
+| Hook | When it fires |
+|------|---------------|
+| `on_train_begin` | Before the first step |
+| `on_train_end` | After the last step |
+| `on_epoch_begin` | Start of each epoch |
+| `on_epoch_end` | End of each epoch |
+| `on_step_begin` | Before each training step |
+| `on_step_end` | After each training step |
+| `on_log` | When metrics are logged |
+| `on_evaluate` | After evaluation |
+| `on_save` | After saving a checkpoint |
+
+## Learning Rate Schedulers
+
+Three schedules are available, all with optional linear warmup:
+
+```ruby
+# Linear warmup then linear decay to 0 (default)
+args = Trainers::TrainingArguments.new(lr_scheduler_type: :linear, warmup_steps: 100)
+
+# Linear warmup then cosine decay to 0
+args = Trainers::TrainingArguments.new(lr_scheduler_type: :cosine, warmup_steps: 100)
+
+# Linear warmup then constant
+args = Trainers::TrainingArguments.new(lr_scheduler_type: :constant, warmup_steps: 100)
+```
+
+## Data Utilities
+
+### Dataset
+
+Wrap an array of hashes:
+
+```ruby
+data = [
+  { input_ids: [101, 2023, 2003], attention_mask: [1, 1, 1], labels: 1 },
+  { input_ids: [101, 2919, 2143], attention_mask: [1, 1, 1], labels: 0 },
+]
+dataset = Trainers::Dataset.new(data)
+```
+
+### Data Collators
+
+Dynamic padding collator (pads each batch to the longest sequence in that batch):
+
+```ruby
+collator = Trainers::DataCollatorWithPadding.new(tokenizer: tokenizer)
+```
+
+Default collator (no padding, expects uniform-length inputs):
+
+```ruby
+collator = Trainers::DefaultDataCollator.new
+```
+
+## Supported Tasks
+
+trainers-rb works with any `Torch::NN::Module`. The `Trainers.from_pretrained` convenience method supports these transformers-rb model classes:
+
+| Task | Model class |
+|------|-------------|
+| `:sequence_classification` | `AutoModelForSequenceClassification` |
+| `:token_classification` | `AutoModelForTokenClassification` |
+| `:question_answering` | `AutoModelForQuestionAnswering` |
+
+You can also use any custom model:
+
+```ruby
+trainer = Trainers::Trainer.new(model: my_custom_model, args: args, ...)
+```
+
+## Device Support
+
+trainers-rb auto-detects the best available device:
+
+- **CPU** — always available
+- **MPS** — Apple Silicon GPU, used automatically when available
+
+```ruby
+# Force CPU
+args = Trainers::TrainingArguments.new(no_mps: true)
+
+# Or set explicitly
+args = Trainers::TrainingArguments.new(device: Torch.device("mps"))
+```
+
+## Architecture
+
+```
+trainers-rb
+  -> transformers-rb    (model loading, tokenizers, HF Hub)
+    -> torch-rb         (autograd, nn modules, optimizers)
+    -> tokenizers        (HuggingFace Rust tokenizers via FFI)
+    -> safetensors       (weight file I/O)
+```
+
+trainers-rb adds the training layer that transformers-rb intentionally omits. Both gems call into the same LibTorch C++ kernels for the actual computation.
+
+## Roadmap
+
+- [ ] More model architectures in transformers-rb (GPT-2, Llama for text generation)
+- [ ] Mixed precision training (fp16/bf16)
+- [ ] Gradient checkpointing for memory efficiency
+- [ ] Dataset streaming for large datasets
+- [ ] Distributed training
+- [ ] Integration with ONNX export for deployment
+- [ ] QLoRA (quantized base model + LoRA)
+
+## Contributing
+
+Bug reports and pull requests are welcome on GitHub.
+
+## License
+
+MIT

data/lib/trainers/callbacks.rb

@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+module Trainers
+  class TrainerState
+    attr_accessor :epoch, :global_step, :max_steps, :num_train_epochs,
+                  :total_flos, :best_metric, :best_model_checkpoint,
+                  :log_history
+
+    def initialize
+      @epoch                 = 0.0
+      @global_step           = 0
+      @max_steps             = 0
+      @num_train_epochs      = 0
+      @total_flos            = 0
+      @best_metric           = nil
+      @best_model_checkpoint = nil
+      @log_history           = []
+    end
+  end
+
+  class TrainerControl
+    attr_accessor :should_training_stop, :should_epoch_stop,
+                  :should_save, :should_evaluate, :should_log
+
+    def initialize
+      @should_training_stop = false
+      @should_epoch_stop    = false
+      @should_save          = false
+      @should_evaluate      = false
+      @should_log           = false
+    end
+  end
+
+  # Base class for trainer callbacks. Override any hook you need.
+  class TrainerCallback
+    def on_train_begin(args, state, control, **kwargs); end
+    def on_train_end(args, state, control, **kwargs); end
+    def on_epoch_begin(args, state, control, **kwargs); end
+    def on_epoch_end(args, state, control, **kwargs); end
+    def on_step_begin(args, state, control, **kwargs); end
+    def on_step_end(args, state, control, **kwargs); end
+    def on_log(args, state, control, logs: nil, **kwargs); end
+    def on_evaluate(args, state, control, metrics: nil, **kwargs); end
+    def on_save(args, state, control, **kwargs); end
+  end
+
+  # Default callback: prints training progress to stdout
+  class PrinterCallback < TrainerCallback
+    def on_log(args, state, control, logs: nil, **kwargs)
+      return unless logs
+      output = logs.map { |k, v| "#{k}: #{format_value(v)}" }.join("  ")
+      puts "[step #{state.global_step}] #{output}"
+    end
+
+    def on_train_begin(args, state, control, **kwargs)
+      puts "Starting training: #{state.num_train_epochs} epochs, #{state.max_steps} total steps"
+    end
+
+    def on_train_end(args, state, control, **kwargs)
+      puts "Training complete. Total steps: #{state.global_step}"
+    end
+
+    def on_evaluate(args, state, control, metrics: nil, **kwargs)
+      return unless metrics
+      output = metrics.map { |k, v| "#{k}: #{format_value(v)}" }.join("  ")
+      puts "[eval step #{state.global_step}] #{output}"
+    end
+
+    private
+
+    def format_value(v)
+      v.is_a?(Float) ? format("%.4f", v) : v.to_s
+    end
+  end
+
+  # Early stopping callback
+  class EarlyStoppingCallback < TrainerCallback
+    def initialize(patience: 3, threshold: 0.0, metric_name: "eval_loss")
+      @patience    = patience
+      @threshold   = threshold
+      @metric_name = metric_name
+      @best_value  = nil
+      @wait_count  = 0
+    end
+
+    def on_evaluate(args, state, control, metrics: nil, **kwargs)
+      return unless metrics
+
+      current = metrics[@metric_name] || metrics[@metric_name.to_sym]
+      return unless current
+
+      if @best_value.nil? || improved?(current, @best_value)
+        @best_value = current
+        @wait_count = 0
+      else
+        @wait_count += 1
+        if @wait_count >= @patience
+          puts "Early stopping triggered after #{@wait_count} evaluations without improvement"
+          control.should_training_stop = true
+        end
+      end
+    end
+
+    private
+
+    def improved?(current, best)
+      if @metric_name.include?("loss")
+        current < best - @threshold
+      else
+        current > best + @threshold
+      end
+    end
+  end
+
+  # Dispatches callback events to all registered callbacks
+  class CallbackHandler
+    def initialize(callbacks)
+      @callbacks = callbacks
+    end
+
+    def fire(event, args, state, control, **kwargs)
+      @callbacks.each do |cb|
+        cb.send(event, args, state, control, **kwargs) if cb.respond_to?(event)
+      end
+      control
+    end
+  end
+end

data/lib/trainers/data/data_collator.rb

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+
+module Trainers
+  class DataCollatorWithPadding
+    attr_reader :tokenizer, :padding, :max_length, :pad_to_multiple_of
+
+    def initialize(tokenizer:, padding: true, max_length: nil, pad_to_multiple_of: nil)
+      @tokenizer          = tokenizer
+      @padding             = padding
+      @max_length          = max_length
+      @pad_to_multiple_of  = pad_to_multiple_of
+    end
+
+    def call(features)
+      return {} if features.empty?
+
+      keys = features.first.keys
+      batch = {}
+
+      keys.each do |key|
+        values = features.map { |f| f[key] }
+
+        if values.first.is_a?(Array)
+          batch[key] = pad_and_stack(key, values)
+        elsif values.first.is_a?(Torch::Tensor)
+          if values.first.dim == 0
+            batch[key] = Torch.stack(values)
+          else
+            batch[key] = pad_and_stack_tensors(key, values)
+          end
+        elsif values.first.is_a?(Integer)
+          batch[key] = Torch.tensor(values, dtype: :int64)
+        elsif values.first.is_a?(Float)
+          batch[key] = Torch.tensor(values, dtype: :float32)
+        else
+          batch[key] = values
+        end
+      end
+
+      batch
+    end
+
+    private
+
+    def pad_and_stack(key, sequences)
+      max_len = compute_max_length(sequences.map(&:length))
+
+      pad_value = if key.to_s.include?("input_id")
+                    pad_token_id
+                  elsif key.to_s.include?("attention_mask")
+                    0
+                  elsif key.to_s.include?("label")
+                    -100
+                  else
+                    0
+                  end
+
+      padded = sequences.map do |seq|
+        padded_seq = seq + [pad_value] * (max_len - seq.length)
+        padded_seq
+      end
+
+      dtype = key.to_s.include?("attention_mask") ? :int64 : :int64
+      Torch.tensor(padded, dtype: dtype)
+    end
+
+    def pad_and_stack_tensors(key, tensors)
+      max_len = compute_max_length(tensors.map { |t| t.size(0) })
+
+      pad_value = if key.to_s.include?("input_id")
+                    pad_token_id
+                  elsif key.to_s.include?("attention_mask")
+                    0
+                  else
+                    0
+                  end
+
+      padded = tensors.map do |t|
+        if t.size(0) < max_len
+          padding = Torch.full([max_len - t.size(0)], pad_value, dtype: t.dtype)
+          Torch.cat([t, padding])
+        else
+          t[0...max_len]
+        end
+      end
+
+      Torch.stack(padded)
+    end
+
+    def compute_max_length(lengths)
+      max_len = @max_length || lengths.max
+      if @pad_to_multiple_of
+        max_len = ((max_len + @pad_to_multiple_of - 1) / @pad_to_multiple_of) * @pad_to_multiple_of
+      end
+      max_len
+    end
+
+    def pad_token_id
+      if @tokenizer.respond_to?(:pad_token_id)
+        @tokenizer.pad_token_id || 0
+      else
+        0
+      end
+    end
+  end
+
+  class DefaultDataCollator
+    def call(features)
+      return {} if features.empty?
+
+      batch = {}
+      features.first.keys.each do |key|
+        values = features.map { |f| f[key] }
+
+        if values.first.is_a?(Torch::Tensor)
+          batch[key] = Torch.stack(values)
+        elsif values.first.is_a?(Integer)
+          batch[key] = Torch.tensor(values, dtype: :int64)
+        elsif values.first.is_a?(Float)
+          batch[key] = Torch.tensor(values, dtype: :float32)
+        elsif values.first.is_a?(Array)
+          batch[key] = Torch.tensor(values)
+        else
+          batch[key] = values
+        end
+      end
+
+      batch
+    end
+  end
+end

data/lib/trainers/data/dataset.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Trainers
+  class Dataset
+    include Enumerable
+
+    attr_reader :data
+
+    def initialize(data)
+      @data = data
+    end
+
+    def [](index)
+      @data[index]
+    end
+
+    def size
+      @data.size
+    end
+    alias_method :length, :size
+
+    def each(&block)
+      @data.each(&block)
+    end
+  end
+end

data/lib/trainers/lora/lora_config.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Trainers
+  class LoraConfig
+    attr_accessor :r, :lora_alpha, :lora_dropout, :target_modules,
+                  :bias, :task_type
+
+    DEFAULTS = {
+      r:              8,
+      lora_alpha:     16,
+      lora_dropout:   0.0,
+      target_modules: ["query", "value"],  # or :all_linear
+      bias:           :none,               # :none, :all, :lora_only
+      task_type:      :sequence_classification
+    }.freeze
+
+    def initialize(**kwargs)
+      DEFAULTS.each do |key, default|
+        value = kwargs.fetch(key, default)
+        instance_variable_set(:"@#{key}", value)
+      end
+    end
+
+    def scaling
+      @lora_alpha.to_f / @r
+    end
+
+    def to_h
+      DEFAULTS.keys.each_with_object({}) do |key, hash|
+        hash[key] = send(key)
+      end
+    end
+  end
+end