fine 0.1.0

data/lib/fine/training/llm_trainer.rb

@@ -0,0 +1,212 @@
+# frozen_string_literal: true
+
+module Fine
+  module Training
+    # Trainer for causal language model fine-tuning
+    class LLMTrainer
+      attr_reader :model, :config, :train_dataset, :val_dataset
+
+      def initialize(model, config, train_dataset:, val_dataset: nil)
+        @model = model
+        @config = config
+        @train_dataset = train_dataset
+        @val_dataset = val_dataset
+        @device = Fine.device
+        @history = []
+      end
+
+      def fit
+        @model.to(@device)
+        @model.train
+
+        optimizer = create_optimizer
+        scheduler = create_scheduler(optimizer)
+
+        @config.callbacks.each { |cb| cb.on_train_begin(model: @model, config: @config) }
+
+        @config.epochs.times do |epoch|
+          epoch_loss = train_epoch(optimizer, scheduler, epoch)
+
+          metrics = { loss: epoch_loss }
+
+          # Validation
+          if @val_dataset
+            val_loss, val_perplexity = evaluate
+            metrics[:val_loss] = val_loss
+            metrics[:val_perplexity] = val_perplexity
+          end
+
+          @history << { epoch: epoch + 1, **metrics }
+
+          @config.callbacks.each { |cb| cb.on_epoch_end(self, epoch + 1, metrics) }
+        end
+
+        @config.callbacks.each { |cb| cb.on_train_end(history: @history) }
+
+        @history
+      end
+
+      def evaluate
+        @model.eval
+        total_loss = 0.0
+        num_batches = 0
+
+        data_loader = create_data_loader(@val_dataset, shuffle: false)
+
+        Torch.no_grad do
+          data_loader.each do |batch|
+            batch = move_to_device(batch)
+
+            outputs = @model.forward(
+              batch[:input_ids],
+              attention_mask: batch[:attention_mask],
+              labels: batch[:labels]
+            )
+
+            total_loss += outputs[:loss].to(:float32).item
+            num_batches += 1
+          end
+        end
+
+        @model.train
+
+        avg_loss = total_loss / num_batches
+        perplexity = Math.exp(avg_loss)
+
+        [avg_loss, perplexity]
+      end
+
+      private
+
+      def train_epoch(optimizer, scheduler, epoch)
+        total_loss = 0.0
+        num_batches = 0
+
+        data_loader = create_data_loader(@train_dataset, shuffle: true)
+        total_steps = data_loader.size
+
+        @config.callbacks.each do |cb|
+          cb.on_epoch_begin(epoch + 1, total_steps: total_steps) if cb.respond_to?(:on_epoch_begin)
+        end
+
+        data_loader.each_with_index do |batch, step|
+          # Move batch to device
+          input_ids = batch[:input_ids].to(@device)
+          attention_mask = batch[:attention_mask].to(@device)
+          labels = batch[:labels].to(@device)
+
+          # Forward pass
+          outputs = @model.forward(input_ids, attention_mask: attention_mask, labels: labels)
+
+          # Get loss value before backward
+          loss_value = outputs[:loss].detach.to(:float32).item
+
+          # Backward pass - scale loss for gradient accumulation
+          scaled_loss = outputs[:loss] / @config.gradient_accumulation_steps
+          scaled_loss.backward
+
+          # CRITICAL: Clear ALL references to free computation graph
+          scaled_loss = nil
+          outputs = nil
+          input_ids = nil
+          attention_mask = nil
+          labels = nil
+          batch = nil
+
+          if (step + 1) % @config.gradient_accumulation_steps == 0
+            # Gradient clipping
+            if @config.max_grad_norm
+              clip_grad_norm(@model.parameters, @config.max_grad_norm)
+            end
+
+            optimizer.step
+            scheduler&.step
+            optimizer.zero_grad
+
+            # Force GC after each optimizer step to free computation graphs
+            GC.start
+          end
+
+          total_loss += loss_value
+          num_batches += 1
+
+          @config.callbacks.each do |cb|
+            cb.on_batch_end(self, step + 1, loss_value) if cb.respond_to?(:on_batch_end)
+          end
+        end
+
+        total_loss / num_batches
+      end
+
+      def create_data_loader(dataset, shuffle:)
+        Datasets::InstructionDataLoader.new(
+          dataset,
+          batch_size: @config.batch_size,
+          shuffle: shuffle,
+          pad_token_id: @config.pad_token_id || 0
+        )
+      end
+
+      def create_optimizer
+        # Separate weight decay for different parameter groups
+        decay_params = []
+        no_decay_params = []
+
+        @model.named_parameters.each do |name, param|
+          next unless param.requires_grad
+
+          if name.include?("bias") || name.include?("layernorm") || name.include?("norm")
+            no_decay_params << param
+          else
+            decay_params << param
+          end
+        end
+
+        param_groups = [
+          { params: decay_params, weight_decay: @config.weight_decay },
+          { params: no_decay_params, weight_decay: 0.0 }
+        ]
+
+        Torch::Optim::AdamW.new(param_groups, lr: @config.learning_rate)
+      end
+
+      def create_scheduler(optimizer)
+        return nil unless @config.warmup_steps && @config.warmup_steps > 0
+
+        # Linear warmup then constant
+        # Note: torch.rb scheduler support may be limited
+        nil
+      end
+
+      def move_to_device(batch)
+        batch.transform_values { |v| v.to(@device) }
+      end
+
+      # Manual gradient clipping implementation
+      def clip_grad_norm(parameters, max_norm)
+        total_norm = 0.0
+
+        parameters.each do |param|
+          next unless param.grad
+
+          # Convert to float32 for .item (bfloat16 not supported)
+          param_norm = param.grad.data.norm(2).to(:float32).item
+          total_norm += param_norm ** 2
+        end
+
+        total_norm = Math.sqrt(total_norm)
+        clip_coef = max_norm / (total_norm + 1e-6)
+
+        if clip_coef < 1.0
+          parameters.each do |param|
+            next unless param.grad
+
+            param.grad.data.mul!(clip_coef)
+          end
+        end
+
+        total_norm
+      end
+    end
+  end
+end

data/lib/fine/training/text_trainer.rb

@@ -0,0 +1,275 @@
+# frozen_string_literal: true
+
+module Fine
+  module Training
+    # Trainer for text classification models
+    class TextTrainer
+      attr_reader :model, :config, :train_loader, :val_loader, :label_map
+      attr_accessor :stop_training
+
+      def initialize(model, config, train_dataset:, val_dataset: nil)
+        @model = model
+        @config = config
+        @stop_training = false
+        @label_map = train_dataset.label_map
+
+        @train_loader = Datasets::TextDataLoader.new(
+          train_dataset,
+          batch_size: config.batch_size,
+          shuffle: true
+        )
+
+        @val_loader = if val_dataset
+          Datasets::TextDataLoader.new(
+            val_dataset,
+            batch_size: config.batch_size,
+            shuffle: false
+          )
+        end
+
+        @history = []
+      end
+
+      # Train the model
+      #
+      # @return [Array<Hash>] Training history
+      def fit
+        @model.train
+
+        optimizer = build_optimizer
+        scheduler = build_scheduler(optimizer)
+
+        run_callbacks(:on_train_begin, self)
+
+        @config.epochs.times do |epoch|
+          break if @stop_training
+
+          run_callbacks(:on_epoch_begin, self, epoch)
+
+          train_metrics = train_epoch(optimizer, epoch)
+
+          val_metrics = @val_loader ? evaluate : {}
+
+          scheduler&.step
+
+          metrics = train_metrics.merge(
+            val_metrics.transform_keys { |k| :"val_#{k}" }
+          )
+          @history << metrics
+
+          run_callbacks(:on_epoch_end, self, epoch, metrics)
+        end
+
+        run_callbacks(:on_train_end, self)
+
+        @history
+      end
+
+      # Evaluate on validation set
+      def evaluate
+        @model.eval
+
+        total_loss = 0.0
+        correct = 0
+        total = 0
+
+        Torch.no_grad do
+          @val_loader.each_batch do |batch|
+            output = @model.call(
+              batch[:input_ids],
+              attention_mask: batch[:attention_mask],
+              token_type_ids: batch[:token_type_ids],
+              labels: batch[:labels]
+            )
+
+            total_loss += output[:loss].item * batch[:labels].size(0)
+            predictions = output[:logits].argmax(dim: 1)
+            correct += predictions.eq(batch[:labels]).sum.item
+            total += batch[:labels].size(0)
+          end
+        end
+
+        @model.train
+
+        {
+          loss: total_loss / total,
+          accuracy: correct.to_f / total
+        }
+      end
+
+      private
+
+      def train_epoch(optimizer, _epoch)
+        total_loss = 0.0
+        correct = 0
+        total = 0
+
+        @train_loader.each_batch.with_index do |batch, batch_idx|
+          run_callbacks(:on_batch_begin, self, batch_idx)
+
+          optimizer.zero_grad
+
+          output = @model.call(
+            batch[:input_ids],
+            attention_mask: batch[:attention_mask],
+            token_type_ids: batch[:token_type_ids],
+            labels: batch[:labels]
+          )
+
+          loss = output[:loss]
+          loss.backward
+          optimizer.step
+
+          batch_loss = loss.item
+          total_loss += batch_loss * batch[:labels].size(0)
+
+          predictions = output[:logits].argmax(dim: 1)
+          correct += predictions.eq(batch[:labels]).sum.item
+          total += batch[:labels].size(0)
+
+          run_callbacks(:on_batch_end, self, batch_idx, batch_loss)
+        end
+
+        {
+          loss: total_loss / total,
+          accuracy: correct.to_f / total
+        }
+      end
+
+      def build_optimizer
+        params = @model.parameters.select(&:requires_grad)
+
+        case @config.optimizer
+        when :adam
+          Torch::Optim::Adam.new(params, lr: @config.learning_rate)
+        when :adamw
+          Torch::Optim::AdamW.new(
+            params,
+            lr: @config.learning_rate,
+            weight_decay: @config.weight_decay
+          )
+        when :sgd
+          Torch::Optim::SGD.new(params, lr: @config.learning_rate, momentum: 0.9)
+        else
+          raise ConfigurationError, "Unknown optimizer: #{@config.optimizer}"
+        end
+      end
+
+      def build_scheduler(optimizer)
+        return nil unless @config.scheduler
+
+        case @config.scheduler
+        when :cosine
+          Torch::Optim::LRScheduler::CosineAnnealingLR.new(optimizer, @config.epochs)
+        when :linear
+          # Linear decay
+          Torch::Optim::LRScheduler::StepLR.new(optimizer, step_size: 1, gamma: 0.9)
+        else
+          nil
+        end
+      end
+
+      def run_callbacks(method, *args)
+        @config.callbacks.each do |callback|
+          callback.send(method, *args)
+        end
+      end
+    end
+
+    # Trainer for text embedding models (contrastive learning)
+    class EmbeddingTrainer
+      attr_reader :model, :config, :train_loader
+      attr_accessor :stop_training
+
+      def initialize(model, config, train_dataset:)
+        @model = model
+        @config = config
+        @stop_training = false
+
+        @train_loader = Datasets::TextDataLoader.new(
+          train_dataset,
+          batch_size: config.batch_size,
+          shuffle: true
+        )
+
+        @history = []
+      end
+
+      def fit
+        @model.train
+
+        optimizer = build_optimizer
+
+        run_callbacks(:on_train_begin, self)
+
+        @config.epochs.times do |epoch|
+          break if @stop_training
+
+          run_callbacks(:on_epoch_begin, self, epoch)
+
+          metrics = train_epoch(optimizer)
+          @history << metrics
+
+          run_callbacks(:on_epoch_end, self, epoch, metrics)
+        end
+
+        run_callbacks(:on_train_end, self)
+
+        @history
+      end
+
+      private
+
+      def train_epoch(optimizer)
+        total_loss = 0.0
+        num_batches = 0
+
+        @train_loader.each_batch do |batch|
+          optimizer.zero_grad
+
+          # For pair datasets, we get anchor and positive texts
+          embeddings = @model.encode(
+            batch[:input_ids],
+            attention_mask: batch[:attention_mask]
+          )
+
+          # Multiple Negatives Ranking Loss
+          # Treat other samples in batch as negatives
+          loss = multiple_negatives_ranking_loss(embeddings)
+
+          loss.backward
+          optimizer.step
+
+          total_loss += loss.item
+          num_batches += 1
+        end
+
+        { loss: total_loss / num_batches }
+      end
+
+      def multiple_negatives_ranking_loss(embeddings, scale: 20.0)
+        # Split embeddings into anchors and positives (assuming paired data)
+        batch_size = embeddings.size(0) / 2
+        anchors = embeddings[0...batch_size]
+        positives = embeddings[batch_size..]
+
+        # Compute similarity matrix
+        scores = Torch.matmul(anchors, positives.transpose(0, 1)) * scale
+
+        # Labels: diagonal is positive (index i matches index i)
+        labels = Torch.arange(batch_size, device: embeddings.device)
+
+        Torch::NN::Functional.cross_entropy(scores, labels)
+      end
+
+      def build_optimizer
+        params = @model.parameters.select(&:requires_grad)
+        Torch::Optim::AdamW.new(params, lr: @config.learning_rate, weight_decay: @config.weight_decay)
+      end
+
+      def run_callbacks(method, *args)
+        @config.callbacks.each { |cb| cb.send(method, *args) }
+      end
+    end
+  end
+end

data/lib/fine/training/trainer.rb

@@ -0,0 +1,194 @@
+# frozen_string_literal: true
+
+module Fine
+  module Training
+    # Main training orchestrator
+    class Trainer
+      attr_reader :model, :config, :train_loader, :val_loader, :label_map
+      attr_accessor :stop_training
+
+      def initialize(model, config, train_dataset:, val_dataset: nil)
+        @model = model
+        @config = config
+        @stop_training = false
+        @label_map = train_dataset.label_map
+
+        # Create data loaders
+        @train_loader = Datasets::DataLoader.new(
+          train_dataset,
+          batch_size: config.batch_size,
+          shuffle: true
+        )
+
+        @val_loader = if val_dataset
+          Datasets::DataLoader.new(
+            val_dataset,
+            batch_size: config.batch_size,
+            shuffle: false
+          )
+        end
+
+        # History tracking
+        @history = []
+      end
+
+      # Train the model
+      #
+      # @return [Array<Hash>] Training history (metrics per epoch)
+      def fit
+        @model.train
+
+        # Build optimizer
+        optimizer = build_optimizer
+        scheduler = build_scheduler(optimizer)
+
+        run_callbacks(:on_train_begin, self)
+
+        @config.epochs.times do |epoch|
+          break if @stop_training
+
+          run_callbacks(:on_epoch_begin, self, epoch)
+
+          # Training epoch
+          train_metrics = train_epoch(optimizer, epoch)
+
+          # Validation
+          val_metrics = @val_loader ? evaluate : {}
+
+          # Step scheduler
+          scheduler&.step
+
+          # Combine metrics
+          metrics = train_metrics.merge(
+            val_metrics.transform_keys { |k| :"val_#{k}" }
+          )
+          @history << metrics
+
+          run_callbacks(:on_epoch_end, self, epoch, metrics)
+        end
+
+        run_callbacks(:on_train_end, self)
+
+        @history
+      end
+
+      # Evaluate on validation set
+      #
+      # @return [Hash] Evaluation metrics
+      def evaluate
+        @model.eval
+
+        total_loss = 0.0
+        correct = 0
+        total = 0
+
+        Torch.no_grad do
+          @val_loader.each_batch do |batch|
+            output = @model.call(batch[:pixel_values], labels: batch[:labels])
+
+            total_loss += output[:loss].item * batch[:labels].size(0)
+            predictions = output[:logits].argmax(dim: 1)
+            correct += (predictions == batch[:labels]).sum.item
+            total += batch[:labels].size(0)
+          end
+        end
+
+        @model.train
+
+        {
+          loss: total_loss / total,
+          accuracy: correct.to_f / total
+        }
+      end
+
+      private
+
+      def train_epoch(optimizer, epoch)
+        total_loss = 0.0
+        correct = 0
+        total = 0
+
+        @train_loader.each_batch.with_index do |batch, batch_idx|
+          run_callbacks(:on_batch_begin, self, batch_idx)
+
+          # Zero gradients
+          optimizer.zero_grad
+
+          # Forward pass
+          output = @model.call(batch[:pixel_values], labels: batch[:labels])
+          loss = output[:loss]
+
+          # Backward pass
+          loss.backward
+
+          # Update weights
+          optimizer.step
+
+          # Track metrics
+          batch_loss = loss.item
+          total_loss += batch_loss * batch[:labels].size(0)
+
+          predictions = output[:logits].argmax(dim: 1)
+          correct += predictions.eq(batch[:labels]).sum.item
+          total += batch[:labels].size(0)
+
+          run_callbacks(:on_batch_end, self, batch_idx, batch_loss)
+        end
+
+        {
+          loss: total_loss / total,
+          accuracy: correct.to_f / total
+        }
+      end
+
+      def build_optimizer
+        params = @model.parameters.select(&:requires_grad)
+
+        case @config.optimizer
+        when :adam
+          Torch::Optim::Adam.new(params, lr: @config.learning_rate)
+        when :adamw
+          Torch::Optim::AdamW.new(
+            params,
+            lr: @config.learning_rate,
+            weight_decay: @config.weight_decay
+          )
+        when :sgd
+          Torch::Optim::SGD.new(
+            params,
+            lr: @config.learning_rate,
+            momentum: 0.9
+          )
+        else
+          raise ConfigurationError, "Unknown optimizer: #{@config.optimizer}"
+        end
+      end
+
+      def build_scheduler(optimizer)
+        return nil unless @config.scheduler
+
+        case @config.scheduler
+        when :cosine
+          Torch::Optim::LRScheduler::CosineAnnealingLR.new(
+            optimizer,
+            @config.epochs
+          )
+        when :step
+          Torch::Optim::LRScheduler::StepLR.new(
+            optimizer,
+            step_size: @config.epochs / 3,
+            gamma: 0.1
+          )
+        else
+          nil
+        end
+      end
+
+      def run_callbacks(method, *args)
+        @config.callbacks.each do |callback|
+          callback.send(method, *args)
+        end
+      end
+    end
+  end
+end

data/lib/fine/transforms/compose.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Fine
+  module Transforms
+    # Composes multiple transforms into a single callable
+    class Compose
+      attr_reader :transforms
+
+      def initialize(transforms)
+        @transforms = transforms
+      end
+
+      def call(image)
+        @transforms.reduce(image) { |img, transform| transform.call(img) }
+      end
+
+      def <<(transform)
+        @transforms << transform
+        self
+      end
+
+      def prepend(transform)
+        @transforms.unshift(transform)
+        self
+      end
+    end
+  end
+end