trainers-rb 0.1.0

data/lib/trainers/trainer.rb

@@ -0,0 +1,340 @@
+# frozen_string_literal: true
+
+module Trainers
+  class Trainer
+    attr_reader :model, :args, :train_dataset, :eval_dataset, :tokenizer,
+                :data_collator, :optimizer, :lr_scheduler, :state, :control
+
+    def initialize(
+      model:,
+      args: nil,
+      train_dataset: nil,
+      eval_dataset: nil,
+      tokenizer: nil,
+      data_collator: nil,
+      compute_metrics: nil,
+      callbacks: []
+    )
+      @model           = model
+      @args            = args || TrainingArguments.new
+      @train_dataset   = train_dataset
+      @eval_dataset    = eval_dataset
+      @tokenizer       = tokenizer
+      @data_collator   = data_collator || DefaultDataCollator.new
+      @compute_metrics = compute_metrics
+      @state           = TrainerState.new
+      @control         = TrainerControl.new
+
+      all_callbacks = [PrinterCallback.new] + callbacks
+      @callback_handler = CallbackHandler.new(all_callbacks)
+    end
+
+    def train
+      device = @args.resolved_device
+      @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
+
+      @state.max_steps       = total_steps
+      @state.num_train_epochs = @args.num_train_epochs
+
+      @optimizer    = create_optimizer
+      @lr_scheduler = create_scheduler(total_steps)
+
+      @callback_handler.fire(:on_train_begin, @args, @state, @control)
+
+      @args.num_train_epochs.times do |epoch|
+        @state.epoch = epoch + 1
+        @callback_handler.fire(:on_epoch_begin, @args, @state, @control)
+        @model.train
+
+        epoch_loss   = 0.0
+        epoch_steps  = 0
+
+        each_batch(@train_dataset, batch_size, shuffle: true) do |batch|
+          @callback_handler.fire(:on_step_begin, @args, @state, @control)
+
+          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
+
+          scaled_loss.backward
+
+          epoch_loss  += loss.item
+          epoch_steps += 1
+          @state.global_step += 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
+
+          # Logging
+          if should_log?
+            logs = {
+              loss:          epoch_loss / epoch_steps,
+              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
+
+          # Step-based evaluation
+          if @args.eval_strategy == :steps && @args.eval_steps &&
+             @state.global_step % @args.eval_steps == 0
+            metrics = evaluate
+            @callback_handler.fire(:on_evaluate, @args, @state, @control, metrics: metrics)
+          end
+
+          # Step-based saving
+          if @args.save_strategy == :steps && @args.save_steps &&
+             @state.global_step % @args.save_steps == 0
+            save_checkpoint
+            @callback_handler.fire(:on_save, @args, @state, @control)
+          end
+
+          @callback_handler.fire(:on_step_end, @args, @state, @control)
+          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-based evaluation
+        if @args.eval_strategy == :epoch && @eval_dataset
+          metrics = evaluate
+          @callback_handler.fire(:on_evaluate, @args, @state, @control, metrics: metrics)
+        end
+
+        # Epoch-based saving
+        if @args.save_strategy == :epoch
+          save_checkpoint
+          @callback_handler.fire(:on_save, @args, @state, @control)
+        end
+
+        @callback_handler.fire(:on_epoch_end, @args, @state, @control)
+        @control.should_epoch_stop = false
+        break if @control.should_training_stop
+      end
+
+      @callback_handler.fire(:on_train_end, @args, @state, @control)
+      @state
+    end
+
+    def evaluate(eval_dataset: nil)
+      dataset = eval_dataset || @eval_dataset
+      raise ArgumentError, "No eval_dataset provided" unless dataset
+
+      device = @args.resolved_device
+      @model.eval
+
+      all_preds  = []
+      all_labels = []
+      total_loss = 0.0
+      total_steps = 0
+
+      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")
+
+          output = forward(batch)
+
+          if labels
+            logits = output.respond_to?(:logits) ? output.logits : output
+            loss = Torch::NN::F.cross_entropy(logits, labels)
+            total_loss += loss.item
+            all_labels << labels.detach.cpu
+          end
+          total_steps += 1
+
+          logits = output.respond_to?(:logits) ? output.logits : output
+          all_preds << logits.detach.cpu
+        end
+      end
+
+      @model.train
+
+      metrics = {}
+      metrics[:eval_loss] = total_loss / total_steps if total_steps > 0
+
+      if @compute_metrics && all_preds.any? && all_labels.any?
+        preds  = Torch.cat(all_preds)
+        labels = Torch.cat(all_labels)
+        eval_pred = EvalPrediction.new(predictions: preds, label_ids: labels)
+        custom_metrics = @compute_metrics.call(eval_pred)
+        metrics.merge!(custom_metrics)
+      end
+
+      metrics
+    end
+
+    def predict(test_dataset)
+      device = @args.resolved_device
+      @model.eval
+
+      all_preds = []
+      Torch.no_grad do
+        each_batch(test_dataset, @args.per_device_eval_batch_size) do |batch|
+          batch  = move_to_device(batch, device)
+          output = forward(batch)
+          logits = output.respond_to?(:logits) ? output.logits : output
+          all_preds << logits.detach.cpu
+        end
+      end
+
+      Torch.cat(all_preds)
+    end
+
+    def save_model(output_dir = nil)
+      output_dir ||= @args.output_dir
+      SaveUtils.save_pretrained(@model, @tokenizer, output_dir, training_args: @args)
+    end
+
+    private
+
+    def compute_loss(batch)
+      labels = batch.delete(:labels) || batch.delete("labels")
+
+      # Try passing labels to the model (some models compute loss internally)
+      output = begin
+        forward(labels ? batch.merge(labels: labels) : batch)
+      rescue => e
+        # If the model doesn't support labels kwarg (e.g. transformers-rb Todo),
+        # fall back to forward without labels + external loss
+        if e.message.include?("Todo") || e.message.include?("not implemented")
+          forward(batch)
+        else
+          raise
+        end
+      end
+
+      # Restore labels to batch for downstream use
+      batch[:labels] = labels if labels
+
+      if output.respond_to?(:loss) && output.loss
+        output.loss
+      elsif labels
+        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. " \
+              "Either pass labels in your dataset or use a model that computes loss."
+      end
+    end
+
+    def forward(batch)
+      if batch.is_a?(Hash)
+        # Filter to only keys the model accepts, using symbol keys
+        @model.call(**batch)
+      else
+        @model.call(batch)
+      end
+    end
+
+    def create_optimizer
+      Optimization.create_optimizer(@model, @args)
+    end
+
+    def create_scheduler(total_steps)
+      warmup_steps = if @args.warmup_steps > 0
+                       @args.warmup_steps
+                     elsif @args.warmup_ratio > 0
+                       (total_steps * @args.warmup_ratio).to_i
+                     else
+                       0
+                     end
+
+      Optimization.create_scheduler(
+        @args.lr_scheduler_type,
+        @optimizer,
+        num_warmup_steps:    warmup_steps,
+        num_training_steps:  total_steps
+      )
+    end
+
+    def each_batch(dataset, batch_size, shuffle: false)
+      indices = (0...dataset.size).to_a
+      indices.shuffle! if shuffle
+
+      (0...dataset.size).step(batch_size) do |start|
+        batch_indices = indices[start, batch_size]
+        next if batch_indices.nil? || batch_indices.empty?
+
+        features = batch_indices.map { |i| dataset[i] }
+        batch = @data_collator.call(features)
+        yield batch
+      end
+    end
+
+    def move_to_device(batch, device)
+      batch.each_with_object({}) do |(key, value), result|
+        result[key] = if value.is_a?(Torch::Tensor)
+                        value.to(device)
+                      else
+                        value
+                      end
+      end
+    end
+
+    def clip_grad_norm!(parameters, max_norm)
+      params = parameters.select { |p| p.grad }
+      return 0.0 if params.empty?
+
+      total_norm_sq = 0.0
+      params.each do |p|
+        total_norm_sq += p.grad.data.norm(2).item ** 2
+      end
+      total_norm = Math.sqrt(total_norm_sq)
+
+      clip_coef = max_norm / (total_norm + 1e-6)
+      if clip_coef < 1.0
+        params.each { |p| p.grad.data.mul!(clip_coef) }
+      end
+
+      total_norm
+    end
+
+    def current_lr
+      @optimizer.param_groups.first[:lr]
+    end
+
+    def should_log?
+      return true if @args.logging_first_step && @state.global_step == 1
+      @state.global_step % @args.logging_steps == 0
+    end
+
+    def save_checkpoint
+      dir = File.join(@args.output_dir, "checkpoint-#{@state.global_step}")
+      save_model(dir)
+      cleanup_checkpoints if @args.save_total_limit
+    end
+
+    def cleanup_checkpoints
+      return unless @args.save_total_limit
+
+      checkpoints = Dir.glob(File.join(@args.output_dir, "checkpoint-*"))
+                       .sort_by { |d| d[/checkpoint-(\d+)/, 1].to_i }
+
+      while checkpoints.length > @args.save_total_limit
+        old = checkpoints.shift
+        FileUtils.rm_rf(old)
+      end
+    end
+  end
+end

data/lib/trainers/trainer_utils.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Trainers
+  module EvalStrategy
+    NO    = :no
+    EPOCH = :epoch
+    STEPS = :steps
+  end
+
+  module SaveStrategy
+    NO    = :no
+    EPOCH = :epoch
+    STEPS = :steps
+  end
+
+  module SchedulerType
+    LINEAR   = :linear
+    COSINE   = :cosine
+    CONSTANT = :constant
+  end
+
+  class EvalPrediction
+    attr_reader :predictions, :label_ids
+
+    def initialize(predictions:, label_ids:)
+      @predictions = predictions
+      @label_ids   = label_ids
+    end
+  end
+end

data/lib/trainers/training_arguments.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module Trainers
+  class TrainingArguments
+    DEFAULTS = {
+      output_dir:                    "./output",
+      num_train_epochs:              3,
+      per_device_train_batch_size:   8,
+      per_device_eval_batch_size:    8,
+      learning_rate:                 5e-5,
+      weight_decay:                  0.0,
+      adam_beta1:                    0.9,
+      adam_beta2:                    0.999,
+      adam_epsilon:                  1e-8,
+      max_grad_norm:                 1.0,
+      gradient_accumulation_steps:   1,
+      warmup_steps:                  0,
+      warmup_ratio:                  0.0,
+      lr_scheduler_type:             :linear,
+      eval_strategy:                 :no,
+      eval_steps:                    nil,
+      save_strategy:                 :epoch,
+      save_steps:                    500,
+      save_total_limit:              nil,
+      logging_steps:                 500,
+      logging_first_step:            false,
+      seed:                          42,
+      device:                        nil,
+      no_mps:                        false,
+      dataloader_drop_last:          false,
+      label_names:                   ["labels"]
+    }.freeze
+
+    DEFAULTS.each_key do |key|
+      attr_accessor key
+    end
+
+    def initialize(**kwargs)
+      DEFAULTS.each do |key, default|
+        value = kwargs.fetch(key, default)
+        instance_variable_set(:"@#{key}", value)
+      end
+
+      unknown = kwargs.keys - DEFAULTS.keys
+      raise ArgumentError, "Unknown arguments: #{unknown.join(', ')}" unless unknown.empty?
+    end
+
+    def resolved_device
+      return @device if @device
+
+      if !@no_mps && defined?(Torch::Backends::MPS) && Torch::Backends::MPS.available?
+        Torch.device("mps")
+      else
+        Torch.device("cpu")
+      end
+    end
+
+    def to_h
+      DEFAULTS.keys.each_with_object({}) do |key, hash|
+        hash[key] = send(key)
+      end
+    end
+  end
+end

data/lib/trainers/version.rb

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

data/lib/trainers-rb.rb

@@ -0,0 +1 @@
+require_relative "trainers"

data/lib/trainers.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "torch"
+require "json"
+require "fileutils"
+
+require_relative "trainers/version"
+require_relative "trainers/trainer_utils"
+require_relative "trainers/training_arguments"
+require_relative "trainers/data/dataset"
+require_relative "trainers/data/data_collator"
+require_relative "trainers/optimization/optimizer"
+require_relative "trainers/optimization/scheduler"
+require_relative "trainers/callbacks"
+require_relative "trainers/save_utils"
+require_relative "trainers/lora/lora_config"
+require_relative "trainers/lora/lora_linear"
+require_relative "trainers/lora/lora_utils"
+require_relative "trainers/lora/lora_model"
+require_relative "trainers/trainer"
+
+module Trainers
+  # Convenience method: load model + tokenizer and prepare for training
+  def self.from_pretrained(model_name, task: :sequence_classification, num_labels: 2)
+    require "transformers-rb"
+
+    model_class = case task
+                  when :sequence_classification
+                    Transformers::AutoModelForSequenceClassification
+                  when :token_classification
+                    Transformers::AutoModelForTokenClassification
+                  when :question_answering
+                    Transformers::AutoModelForQuestionAnswering
+                  else
+                    Transformers::AutoModel
+                  end
+
+    model     = model_class.from_pretrained(model_name, num_labels: num_labels)
+    tokenizer = Transformers::AutoTokenizer.from_pretrained(model_name)
+
+    [model, tokenizer]
+  end
+end

metadata

@@ -0,0 +1,149 @@
+--- !ruby/object:Gem::Specification
+name: trainers-rb
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Vishwajeetsingh Desurkar
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: torch-rb
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.17.1
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.17.1
+- !ruby/object:Gem::Dependency
+  name: transformers-rb
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.2.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.2.0
+- !ruby/object:Gem::Dependency
+  name: safetensors
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.1.1
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.1.1
+- !ruby/object:Gem::Dependency
+  name: tokenizers
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.5.3
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.5.3
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+- !ruby/object:Gem::Dependency
+  name: minitest
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+description: Training loop, LoRA, and optimization utilities for fine-tuning HuggingFace
+  transformer models using torch-rb and transformers-rb. Supports full fine-tuning,
+  LoRA adapters, learning rate scheduling, callbacks, and model serialization via
+  safetensors.
+email:
+- selectus2@users.noreply.rubygems.org
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- LICENSE.txt
+- README.md
+- lib/trainers-rb.rb
+- lib/trainers.rb
+- lib/trainers/callbacks.rb
+- lib/trainers/data/data_collator.rb
+- lib/trainers/data/dataset.rb
+- lib/trainers/lora/lora_config.rb
+- lib/trainers/lora/lora_linear.rb
+- lib/trainers/lora/lora_model.rb
+- lib/trainers/lora/lora_utils.rb
+- lib/trainers/optimization/optimizer.rb
+- lib/trainers/optimization/scheduler.rb
+- lib/trainers/save_utils.rb
+- lib/trainers/trainer.rb
+- lib/trainers/trainer_utils.rb
+- lib/trainers/training_arguments.rb
+- lib/trainers/version.rb
+homepage: https://github.com/trainers-rb/trainers-rb
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/trainers-rb/trainers-rb
+  source_code_uri: https://github.com/trainers-rb/trainers-rb
+  changelog_uri: https://github.com/trainers-rb/trainers-rb/blob/main/CHANGELOG.md
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.1.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.5
+specification_version: 4
+summary: Fine-tune transformer models in Ruby
+test_files: []