trainers-rb 0.1.2 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b2025ca407a70ba9d0a206306f48486eaae88854e767dd006aa9f37af752eaf1
-  data.tar.gz: 5e4a9e5a90d377791c5577399b88aab55e5d00468f4cae13710e0d140d482e5f
+  metadata.gz: 8db3543c735b87abb79d2de2d8278f11c3999289f2c68929a6cc7c855568b6a1
+  data.tar.gz: ff4783fde9226bff71d59e00960997109a5426f407572ea41859657a120e907f
 SHA512:
-  metadata.gz: 16679778184715a50d08ff4de4e5050bca75c98dc02b542331f14aae79a6a0b857b3a2699da060e3de3ce7924357486b3b60da2563e7759eccd4e2fd10d6c716
-  data.tar.gz: d1580d3571d6adaa1dd771b4cfdbea1bf9ea3a002b35bb06e5f4f9184946a6243dd9aac19c4ac96fccb19690e1e7c2291d51a0c03ad551f7ca32600a7092cc91
+  metadata.gz: 492022d310c9dc8b59e823e550769bdabbbd6af100d503ca82421c4fc3cef3aaea041f667c21d96d223b19fd5441c5b85fbd515a44f4ccb0a6820b471adcd28c
+  data.tar.gz: 92678a59b82c31a055c3e28bfe40b60a6c830368c9707d00ccb2481c43213127e8210d7547d2eff32f75d636c48632b8ef44056e9c0692841cc983b983c5297b

data/CHANGELOG.md

@@ -1,5 +1,33 @@
 # Changelog
 
+## 0.2.0 (2026-10-31)
+
+**This release changes training results.** Runs that use `gradient_accumulation_steps > 1` now follow the intended learning-rate schedule. All runs now apply `seed:`, so data order is different from 0.1.x and reproducible. Expect different (better) final models than 0.1.x with the same arguments.
+
+### Changed
+
+- **A step is now one optimizer update** (#12). `global_step`, `max_steps`, `logging_steps`, `eval_steps`, `save_steps`, `warmup_steps` and the scheduler all count optimizer updates, not batches. The scheduler is sized as `ceil(batches_per_epoch / gradient_accumulation_steps) * num_train_epochs`, so it now decays to the end. Previously, with accumulation 4, a linear schedule ended near 75% of the peak learning rate.
+  - The last, shorter accumulation window of an epoch now makes an optimizer update, with its loss averaged over the batches it holds. Previously its gradients leaked into the next epoch.
+  - `on_step_begin` fires at the start of each accumulation window and `on_step_end` after each optimizer update. Nothing changes with `gradient_accumulation_steps: 1`.
+- **`seed:` is applied** (#13). `train` seeds Torch, Ruby's RNG and its own shuffling. `seed: nil` disables this. New `Trainers.set_seed(seed)` seeds model initialisation.
+- **Progress goes to `Trainers.logger`** on `$stderr`, not stdout (#15). Set `Trainers.logger = nil` for a silent library. The default callback is now `Trainers::LoggingCallback`; `PrinterCallback` still prints to stdout if you add it to `callbacks:`.
+- **Errors descend from `Trainers::Error`** (#16): `ConfigurationError`, `AdapterError` and `CheckpointError`. Code that rescued `RuntimeError` for these must rescue `Trainers::Error` instead.
+- `SaveUtils.save_lora_adapters` and `LoraModel.save_adapters` raise `AdapterError` when the model has no adapters, instead of printing and returning `nil` (#16).
+- **Requires Ruby 3.3 or newer.** Ruby 3.1 and 3.2 are end-of-life.
+- New runtime dependency: `logger`, which Ruby 3.5+ no longer ships as a default gem.
+
+### Fixed
+
+- `dataloader_drop_last: true` skips the final partial batch in training, evaluation and prediction (#14). `train` raises `ConfigurationError` if that leaves no batches.
+- `label_names:` is honoured (#14). Every listed key is passed to the model, and the first is used for the built-in cross-entropy loss and `evaluate`.
+- `LoraModel.apply` works on models whose layers sit in a `ModuleList`, which covers every HuggingFace transformer, for example DistilBERT's `transformer.layer.0.attention.q_lin`. It used to raise `Could not find module part '0'`, because `ModuleList` stores its children in a Hash keyed by `"0"`, not in an Array.
+- When an epoch's last optimizer step was already logged, the end-of-epoch log no longer repeats it.
+- Learning rates and other floats below 0.001 are logged in scientific notation (`2.000e-05`), not as `0.0000`.
+
+### Added
+
+- `TrainingArguments` validates `seed` (nil or a non-negative Integer) and `label_names` (an Array of Strings or Symbols).
+
 ## 0.1.2 (2026-10-30)
 
 ### Fixed

data/README.md

@@ -134,7 +134,7 @@ Trainers::LoraModel.load_adapters(model, "./lora-adapters")
 | `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 |
+| `gradient_accumulation_steps` | `1` | Accumulate gradients over N batches per optimizer step |
 | `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` |
@@ -142,12 +142,60 @@ Trainers::LoraModel.load_adapters(model, "./lora-adapters")
 | `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 |
+| `save_steps` | `500` | Save a checkpoint every N steps (when `save_strategy: :steps`) |
 | `logging_steps` | `500` | Log every N steps |
-| `seed` | `42` | Random seed |
+| `seed` | `42` | Seeds Torch, Ruby's RNG and data shuffling at the start of `train`; `nil` disables seeding |
+| `dataloader_drop_last` | `false` | Skip the final batch when it is smaller than the batch size |
+| `label_names` | `["labels"]` | Batch keys that hold labels; the first is used for the built-in cross-entropy loss |
 | `no_mps` | `false` | Force CPU even if MPS is available |
 
+**A step is one optimizer update.** With `gradient_accumulation_steps: 4`, one step covers four batches, and `max_steps`, `global_step`, `logging_steps`, `eval_steps`, `save_steps`, `warmup_steps` and the learning-rate schedule all count optimizer updates. When an epoch's batch count is not a multiple of the accumulation steps, the last, shorter window still makes an update, with its loss averaged over the batches it contains, so no gradient carries into the next epoch.
+
 Invalid values, such as a negative learning rate, a zero batch size or an unknown `lr_scheduler_type`, raise `ArgumentError` when `TrainingArguments` is created.
 
+## Reproducibility
+
+`Trainer#train` seeds Torch, Ruby's global RNG and its own shuffling from `args.seed`. The model's initial weights are created before training starts, so seed those yourself:
+
+```ruby
+Trainers.set_seed(42)
+model, tokenizer = Trainers.from_pretrained("distilbert-base-uncased", num_labels: 2)
+```
+
+Two runs with the same seed, data and arguments produce identical loss curves on CPU. On Apple Silicon (MPS), some kernels are not deterministic, so expect small differences between runs.
+
+## Logging
+
+trainers-rb reports progress through `Trainers.logger`, which writes to `$stderr` at `INFO` by default. Nothing is written to stdout.
+
+```ruby
+Trainers.logger = Rails.logger                  # use your app's logger
+Trainers.logger.level = Logger::WARN            # only warnings and errors
+Trainers.logger = nil                           # completely silent
+```
+
+To print progress to stdout instead, add `Trainers::PrinterCallback.new` to `callbacks:`.
+
+## Errors
+
+Every error trainers-rb raises descends from `Trainers::Error`:
+
+| Error | Raised when |
+|---|---|
+| `Trainers::ConfigurationError` | The model has no trainable parameters, returns no loss and no labels are present, or the training data yields no batches |
+| `Trainers::AdapterError` | No module matches `target_modules`, an adapter file is missing, or `save_adapters` finds no adapters |
+| `Trainers::CheckpointError` | Reserved for checkpoint save and load errors |
+
+```ruby
+begin
+  trainer.train
+rescue Trainers::Error => e
+  # anything trainers-rb raised
+end
+```
+
+Invalid `TrainingArguments` values raise `ArgumentError`.
+
 ## Callbacks
 
 Built-in callbacks:
@@ -286,11 +334,6 @@ trainers-rb adds the training layer that transformers-rb intentionally omits. Bo
 
 These are open correctness issues. They are tracked, and listed here so nobody discovers them the hard way:
 
-- **The learning-rate schedule is wrong with `gradient_accumulation_steps > 1`.** The scheduler is sized in micro-batches but stepped once per optimizer step, so it never decays fully. ([#12](https://github.com/Selectus2/trainers-rb/issues/12))
-- **`seed:` is never applied.** Runs are not reproducible. ([#13](https://github.com/Selectus2/trainers-rb/issues/13))
-- **`dataloader_drop_last:` and `label_names:` are accepted and ignored.** ([#14](https://github.com/Selectus2/trainers-rb/issues/14))
-- **Progress output goes to stdout through `puts`.** There is no logger to configure or silence. ([#15](https://github.com/Selectus2/trainers-rb/issues/15))
-- **`save_lora_adapters` returns `nil` instead of raising when a model has no adapters.** ([#16](https://github.com/Selectus2/trainers-rb/issues/16))
 - **A saved model cannot be loaded back.** No `config.json` is written. ([#17](https://github.com/Selectus2/trainers-rb/issues/17))
 - **Checkpoints cannot resume training.** Optimizer, scheduler and trainer state are not saved. ([#18](https://github.com/Selectus2/trainers-rb/issues/18))
 - **LoRA checkpoints save the full state dict** rather than just the adapters. ([#19](https://github.com/Selectus2/trainers-rb/issues/19))

data/lib/trainers/callbacks.rb

@@ -44,32 +44,48 @@ module Trainers
     def on_save(args, state, control, **kwargs); end
   end
 
-  # Default callback: prints training progress to stdout
+  # Prints training progress to stdout. Opt in with callbacks: [PrinterCallback.new]
   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}"
+      emit "[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"
+      emit "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}"
+      emit "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}"
+      emit "[eval step #{state.global_step}] #{output}"
     end
 
     private
 
+    def emit(message)
+      puts message
+    end
+
+    # Small floats such as learning rates (2e-5) would print as 0.0000
     def format_value(v)
-      v.is_a?(Float) ? format("%.4f", v) : v.to_s
+      return v.to_s unless v.is_a?(Float)
+
+      v.nonzero? && v.abs < 1e-3 ? format("%.3e", v) : format("%.4f", v)
+    end
+  end
+
+  # Default callback: reports the same progress through Trainers.logger
+  class LoggingCallback < PrinterCallback
+    private
+
+    def emit(message)
+      Trainers.log(:info, message)
     end
   end
 
@@ -95,7 +111,7 @@ module Trainers
       else
         @wait_count += 1
         if @wait_count >= @patience
-          puts "Early stopping triggered after #{@wait_count} evaluations without improvement"
+          Trainers.log(:info, "Early stopping triggered after #{@wait_count} evaluations without improvement")
           control.should_training_stop = true
         end
       end

data/lib/trainers/errors.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Trainers
+  # Base class for every error trainers-rb raises. Rescue this to catch them all.
+  class Error < StandardError; end
+
+  # The model, data or arguments cannot be trained as configured
+  class ConfigurationError < Error; end
+
+  # LoRA adapters are missing, cannot be applied, or cannot be loaded
+  class AdapterError < Error; end
+
+  # A checkpoint cannot be written or read
+  class CheckpointError < Error; end
+end

data/lib/trainers/logging.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "logger"
+
+module Trainers
+  class << self
+    # Where trainers-rb reports progress. Defaults to $stderr at INFO.
+    # Set to any Logger-compatible object, or nil for a silent library.
+    attr_accessor :logger
+
+    def default_logger
+      Logger.new($stderr, level: Logger::INFO, progname: "trainers").tap do |logger|
+        logger.formatter = proc do |severity, _time, progname, message|
+          prefix = severity == "INFO" ? "" : "#{severity}: "
+          "[#{progname}] #{prefix}#{message}\n"
+        end
+      end
+    end
+
+    # Internal: log through Trainers.logger when one is set
+    def log(level, message)
+      logger&.public_send(level, message)
+    end
+  end
+
+  self.logger = default_logger
+end

data/lib/trainers/lora/lora_model.rb

@@ -20,7 +20,7 @@ module Trainers
       targets = LoraUtils.find_target_modules(model, config.target_modules)
 
       if targets.empty?
-        raise "No Linear modules found matching target_modules: #{config.target_modules.inspect}. " \
+        raise AdapterError, "No Linear modules found matching target_modules: #{config.target_modules.inspect}. " \
               "Available modules: #{model.named_modules.map(&:first).join(', ')}"
       end
 
@@ -50,7 +50,7 @@ module Trainers
       end
       # :none — biases stay frozen (default)
 
-      puts "LoRA applied to #{targets.size} modules: #{targets.keys.join(', ')}"
+      Trainers.log(:info, "LoRA applied to #{targets.size} modules: #{targets.keys.join(', ')}")
       LoraUtils.print_trainable_parameters(model)
 
       model
@@ -63,7 +63,7 @@ module Trainers
       lora_modules.each do |name, mod|
         LoraUtils.replace_module(model, name, mod.to_linear)
       end
-      puts "Merged #{lora_modules.size} LoRA adapters into base model"
+      Trainers.log(:info, "Merged #{lora_modules.size} LoRA adapters into base model")
       model
     end
 

data/lib/trainers/lora/lora_utils.rb

@@ -30,24 +30,28 @@ module Trainers
 
       # Navigate to the parent module
       parts[0...-1].each do |part|
-        if part =~ /\A\d+\z/
-          # Numeric index — likely a ModuleList element
-          parent = parent.instance_variable_get(:@modules)[part.to_i]
-        else
-          parent = parent.instance_variable_get(:"@#{part}")
-        end
-
-        raise "Could not find module part '#{part}' in #{target_name}" unless parent
+        parent = child_module(parent, part)
+        raise AdapterError, "Could not find module part '#{part}' in #{target_name}" unless parent
       end
 
       child_name = parts.last
-      if child_name =~ /\A\d+\z/
-        parent.instance_variable_get(:@modules)[child_name.to_i] = new_module
+      registered = parent.instance_variable_get(:@modules)
+      if registered&.key?(child_name)
+        registered[child_name] = new_module
       else
         parent.instance_variable_set(:"@#{child_name}", new_module)
       end
     end
 
+    # ModuleList and add_module children live in @modules, a Hash keyed by
+    # name ("0", "1", ...); everything else is an instance variable.
+    def self.child_module(parent, name)
+      registered = parent.instance_variable_get(:@modules)
+      return registered[name] if registered&.key?(name)
+
+      parent.instance_variable_get(:"@#{name}") if name.match?(/\A[a-z_]\w*\z/i)
+    end
+
     # Count total and trainable parameters
     def self.print_trainable_parameters(model)
       total    = 0
@@ -59,9 +63,9 @@ module Trainers
       end
 
       pct = total > 0 ? (trainable.to_f / total * 100) : 0.0
-      puts "trainable params: #{format_number(trainable)} || " \
-           "all params: #{format_number(total)} || " \
-           "trainable%: #{format('%.4f', pct)}%"
+      Trainers.log(:info, "trainable params: #{format_number(trainable)} || " \
+                          "all params: #{format_number(total)} || " \
+                          "trainable%: #{format('%.4f', pct)}%")
 
       { total: total, trainable: trainable, percentage: pct }
     end

data/lib/trainers/optimization/optimizer.rb

@@ -35,7 +35,7 @@ module Trainers
       end
 
       if param_groups.empty?
-        raise "No trainable parameters found. Did you forget to unfreeze the model or apply LoRA?"
+        raise ConfigurationError, "No trainable parameters found. Did you forget to unfreeze the model or apply LoRA?"
       end
 
       Torch::Optim::AdamW.new(

data/lib/trainers/save_utils.rb

@@ -18,7 +18,7 @@ module Trainers
       model_path = File.join(output_dir, "model.safetensors")
       require "safetensors"
       Safetensors::Torch.save_file(tensors, model_path)
-      puts "Model saved to #{model_path}"
+      Trainers.log(:info, "Model saved to #{model_path}")
 
       # Save tokenizer
       if tokenizer
@@ -49,21 +49,20 @@ module Trainers
       end
 
       if lora_state.empty?
-        puts "No LoRA adapters found in model"
-        return
+        raise AdapterError, "No LoRA adapters found in model. Apply them with LoraModel.apply before saving."
       end
 
       adapter_path = File.join(output_dir, "adapter_model.safetensors")
       require "safetensors"
       Safetensors::Torch.save_file(lora_state, adapter_path)
-      puts "LoRA adapters saved to #{adapter_path} (#{lora_state.size} tensors)"
+      Trainers.log(:info, "LoRA adapters saved to #{adapter_path} (#{lora_state.size} tensors)")
 
       output_dir
     end
 
     def self.load_lora_adapters(model, input_dir)
       adapter_path = File.join(input_dir, "adapter_model.safetensors")
-      raise "Adapter file not found: #{adapter_path}" unless File.exist?(adapter_path)
+      raise AdapterError, "Adapter file not found: #{adapter_path}" unless File.exist?(adapter_path)
 
       require "safetensors"
       lora_state = Safetensors::Torch.load_file(adapter_path)
@@ -78,7 +77,7 @@ module Trainers
         end
       end
 
-      puts "LoRA adapters loaded from #{adapter_path}"
+      Trainers.log(:info, "LoRA adapters loaded from #{adapter_path}")
     end
   end
 end

data/lib/trainers/trainer.rb

@@ -25,25 +25,35 @@ module Trainers
       @state           = TrainerState.new
       @control         = TrainerControl.new
 
-      all_callbacks = [PrinterCallback.new] + callbacks
+      all_callbacks = [LoggingCallback.new] + callbacks
       @callback_handler = CallbackHandler.new(all_callbacks)
     end
 
     def train
       device = @args.resolved_device
+      apply_seed
       @model.to(device)
       @model.train
 
-      num_examples   = @train_dataset.size
-      batch_size     = @args.per_device_train_batch_size
-      steps_per_epoch = (num_examples.to_f / batch_size).ceil
-      total_steps    = steps_per_epoch * @args.num_train_epochs
+      batch_size        = @args.per_device_train_batch_size
+      accum             = @args.gradient_accumulation_steps
+      batches_per_epoch = num_batches(@train_dataset, batch_size)
+      if batches_per_epoch.zero?
+        raise ConfigurationError, "train_dataset yields no batches: #{@train_dataset.size} examples, " \
+                                  "batch size #{batch_size}, dataloader_drop_last #{@args.dataloader_drop_last}"
+      end
+
+      # One step is one optimizer update. The last accumulation window of an
+      # epoch may be shorter than gradient_accumulation_steps; it still steps.
+      steps_per_epoch = (batches_per_epoch.to_f / accum).ceil
+      total_steps     = steps_per_epoch * @args.num_train_epochs
 
-      @state.max_steps       = total_steps
+      @state.max_steps        = total_steps
       @state.num_train_epochs = @args.num_train_epochs
 
       @optimizer    = create_optimizer
       @lr_scheduler = create_scheduler(total_steps)
+      @optimizer.zero_grad
 
       @callback_handler.fire(:on_train_begin, @args, @state, @control)
 
@@ -52,38 +62,38 @@ module Trainers
         @callback_handler.fire(:on_epoch_begin, @args, @state, @control)
         @model.train
 
-        epoch_loss   = 0.0
-        epoch_steps  = 0
+        epoch_loss    = 0.0
+        epoch_batches = 0
+
+        each_batch(@train_dataset, batch_size, shuffle: true).with_index do |batch, index|
+          position     = index % accum
+          window_start = index - position
+          window_size  = [accum, batches_per_epoch - window_start].min
 
-        each_batch(@train_dataset, batch_size, shuffle: true) do |batch|
-          @callback_handler.fire(:on_step_begin, @args, @state, @control)
+          @callback_handler.fire(:on_step_begin, @args, @state, @control) if position.zero?
 
           batch = move_to_device(batch, device)
           loss  = compute_loss(batch)
 
-          scaled_loss = if @args.gradient_accumulation_steps > 1
-                          loss / @args.gradient_accumulation_steps
-                        else
-                          loss
-                        end
-
+          # Average the gradient over the micro-batches in this window
+          scaled_loss = window_size > 1 ? loss / window_size : loss
           scaled_loss.backward
 
-          epoch_loss  += loss.item
-          epoch_steps += 1
-          @state.global_step += 1
+          epoch_loss    += loss.item
+          epoch_batches += 1
 
-          if @state.global_step % @args.gradient_accumulation_steps == 0
-            clip_grad_norm!(@model.parameters, @args.max_grad_norm)
-            @optimizer.step
-            @lr_scheduler.step
-            @optimizer.zero_grad
-          end
+          next unless position == window_size - 1
+
+          clip_grad_norm!(@model.parameters, @args.max_grad_norm)
+          @optimizer.step
+          @lr_scheduler.step
+          @optimizer.zero_grad
+          @state.global_step += 1
 
           # Logging
           if should_log?
             logs = {
-              loss:          epoch_loss / epoch_steps,
+              loss:          epoch_loss / epoch_batches,
               learning_rate: current_lr,
               epoch:         @state.epoch
             }
@@ -109,11 +119,13 @@ module Trainers
           break if @control.should_training_stop || @control.should_epoch_stop
         end
 
-        # Epoch-level logging
-        epoch_avg_loss = epoch_steps > 0 ? epoch_loss / epoch_steps : 0.0
-        logs = { loss: epoch_avg_loss, learning_rate: current_lr, epoch: @state.epoch }
-        @state.log_history << logs.merge(step: @state.global_step)
-        @callback_handler.fire(:on_log, @args, @state, @control, logs: logs)
+        # Epoch-level logging, unless the last step of the epoch already logged
+        unless @state.log_history.last&.fetch(:step, nil) == @state.global_step
+          epoch_avg_loss = epoch_batches > 0 ? epoch_loss / epoch_batches : 0.0
+          logs = { loss: epoch_avg_loss, learning_rate: current_lr, epoch: @state.epoch }
+          @state.log_history << logs.merge(step: @state.global_step)
+          @callback_handler.fire(:on_log, @args, @state, @control, logs: logs)
+        end
 
         # Epoch-based evaluation
         if @args.eval_strategy == :epoch && @eval_dataset
@@ -151,7 +163,7 @@ module Trainers
       Torch.no_grad do
         each_batch(dataset, @args.per_device_eval_batch_size) do |batch|
           batch  = move_to_device(batch, device)
-          labels = batch.delete(:labels) || batch.delete("labels")
+          labels = pop_labels(batch).values.first
 
           output = forward(batch)
 
@@ -209,11 +221,12 @@ module Trainers
     private
 
     def compute_loss(batch)
-      labels = batch.delete(:labels) || batch.delete("labels")
+      label_values = pop_labels(batch)
+      labels       = label_values.values.first
 
       # Try passing labels to the model (some models compute loss internally)
       output = begin
-        forward(labels ? batch.merge(labels: labels) : batch)
+        forward(batch.merge(label_values))
       rescue => e
         # If the model doesn't support labels kwarg (e.g. transformers-rb Todo),
         # fall back to forward without labels + external loss
@@ -225,7 +238,7 @@ module Trainers
       end
 
       # Restore labels to batch for downstream use
-      batch[:labels] = labels if labels
+      batch.merge!(label_values)
 
       if output.respond_to?(:loss) && output.loss
         output.loss
@@ -233,7 +246,7 @@ module Trainers
         logits = output.respond_to?(:logits) ? output.logits : output
         Torch::NN::F.cross_entropy(logits, labels)
       else
-        raise "Model did not return a loss and no labels found in batch. " \
+        raise ConfigurationError, "Model did not return a loss and no labels found in batch. " \
               "Either pass labels in your dataset or use a model that computes loss."
       end
     end
@@ -268,20 +281,40 @@ module Trainers
       )
     end
 
+    # Removes every key named in args.label_names (symbol or string) from the
+    # batch and returns them as { name_symbol => tensor }, in label_names order
+    def pop_labels(batch)
+      @args.label_names.each_with_object({}) do |name, labels|
+        value = batch.key?(name.to_sym) ? batch.delete(name.to_sym) : batch.delete(name.to_s)
+        labels[name.to_sym] = value unless value.nil?
+      end
+    end
+
+    def num_batches(dataset, batch_size)
+      full, rest = dataset.size.divmod(batch_size)
+      rest.positive? && !@args.dataloader_drop_last ? full + 1 : full
+    end
+
     def each_batch(dataset, batch_size, shuffle: false)
-      indices = (0...dataset.size).to_a
-      indices.shuffle! if shuffle
+      return enum_for(:each_batch, dataset, batch_size, shuffle: shuffle) unless block_given?
 
-      (0...dataset.size).step(batch_size) do |start|
-        batch_indices = indices[start, batch_size]
-        next if batch_indices.nil? || batch_indices.empty?
+      indices = (0...dataset.size).to_a
+      indices.shuffle!(random: @rng || Random.new) if shuffle
 
-        features = batch_indices.map { |i| dataset[i] }
-        batch = @data_collator.call(features)
-        yield batch
+      num_batches(dataset, batch_size).times do |i|
+        batch_indices = indices[i * batch_size, batch_size]
+        features = batch_indices.map { |j| dataset[j] }
+        yield @data_collator.call(features)
       end
     end
 
+    # Seeds Torch, Kernel#rand and the shuffle RNG from args.seed (nil skips)
+    def apply_seed
+      seed = @args.seed
+      @rng = seed.nil? ? Random.new : Random.new(seed)
+      Trainers.set_seed(seed) unless seed.nil?
+    end
+
     def move_to_device(batch, device)
       batch.each_with_object({}) do |(key, value), result|
         result[key] = if value.is_a?(Torch::Tensor)

data/lib/trainers/training_arguments.rb

@@ -88,6 +88,8 @@ module Trainers
       check(:warmup_steps, "a non-negative Integer") { |v| v.is_a?(Integer) && !v.negative? }
       check(:warmup_ratio, "a number between 0 and 1") { |v| v.is_a?(Numeric) && v.between?(0, 1) }
       check(:eval_steps, "nil or a positive Integer") { |v| v.nil? || (v.is_a?(Integer) && v.positive?) }
+      check(:seed, "nil or a non-negative Integer") { |v| v.nil? || (v.is_a?(Integer) && !v.negative?) }
+      check(:label_names, "an Array of Strings or Symbols") { |v| v.is_a?(Array) && v.all? { |n| n.is_a?(String) || n.is_a?(Symbol) } }
 
       check(:lr_scheduler_type, "one of #{SCHEDULER_TYPES.map(&:inspect).join(', ')}") { |v| SCHEDULER_TYPES.include?(v) }
       check(:eval_strategy, "one of #{STRATEGIES.map(&:inspect).join(', ')}") { |v| STRATEGIES.include?(v) }

data/lib/trainers/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Trainers
-  VERSION = "0.1.2"
+  VERSION = "0.2.0"
 end

data/lib/trainers.rb

@@ -5,6 +5,8 @@ require "json"
 require "fileutils"
 
 require_relative "trainers/version"
+require_relative "trainers/errors"
+require_relative "trainers/logging"
 require_relative "trainers/trainer_utils"
 require_relative "trainers/training_arguments"
 require_relative "trainers/data/dataset"
@@ -20,6 +22,14 @@ require_relative "trainers/lora/lora_model"
 require_relative "trainers/trainer"
 
 module Trainers
+  # Seeds Torch and Ruby's global RNG. Call before building a model to make
+  # its initial weights reproducible; Trainer#train also calls it with args.seed.
+  def self.set_seed(seed)
+    Torch.manual_seed(seed)
+    srand(seed)
+    seed
+  end
+
   # Convenience method: load model + tokenizer and prepare for training
   def self.from_pretrained(model_name, task: :sequence_classification, num_labels: 2)
     require "transformers-rb"

metadata

@@ -1,14 +1,13 @@
 --- !ruby/object:Gem::Specification
 name: trainers-rb
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.2.0
 platform: ruby
 authors:
 - Vishwajeetsingh Desurkar
-autorequire:
 bindir: bin
 cert_chain: []
-date: 2026-10-30 00:00:00.000000000 Z
+date: 1980-01-02 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: torch-rb
@@ -66,6 +65,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: 0.5.3
+- !ruby/object:Gem::Dependency
+  name: logger
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -112,6 +125,8 @@ files:
 - lib/trainers/callbacks.rb
 - lib/trainers/data/data_collator.rb
 - lib/trainers/data/dataset.rb
+- lib/trainers/errors.rb
+- lib/trainers/logging.rb
 - lib/trainers/lora/lora_config.rb
 - lib/trainers/lora/lora_linear.rb
 - lib/trainers/lora/lora_model.rb
@@ -132,7 +147,6 @@ metadata:
   changelog_uri: https://github.com/Selectus2/trainers-rb/blob/main/CHANGELOG.md
   bug_tracker_uri: https://github.com/Selectus2/trainers-rb/issues
   rubygems_mfa_required: 'true'
-post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -140,15 +154,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: 3.1.0
+      version: 3.3.0
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.5.22
-signing_key:
+rubygems_version: 4.0.20
 specification_version: 4
 summary: Fine-tune transformer models in Ruby
 test_files: []