fine 0.1.0

data/lib/fine/models/sentence_transformer.rb

@@ -0,0 +1,202 @@
+# frozen_string_literal: true
+
+module Fine
+  module Models
+    # Sentence Transformer for generating text embeddings
+    #
+    # Produces dense vector representations of sentences for semantic similarity,
+    # clustering, and retrieval tasks.
+    class SentenceTransformer < Base
+      attr_reader :encoder, :pooling_mode
+
+      POOLING_MODES = %i[cls mean max].freeze
+
+      # Load a pretrained sentence transformer
+      #
+      # @param model_id [String] HuggingFace model ID
+      # @param pooling_mode [Symbol] Pooling strategy (:cls, :mean, :max)
+      # @return [SentenceTransformer]
+      def self.from_pretrained(model_id, pooling_mode: :mean)
+        downloader = Hub::ModelDownloader.new(model_id)
+        model_path = downloader.download
+
+        config = Hub::ConfigLoader.from_pretrained(model_path)
+
+        model = new(config, pooling_mode: pooling_mode)
+
+        # Load weights
+        weights_path = downloader.file_path("model.safetensors")
+        if File.exist?(weights_path)
+          load_pretrained_weights(model, weights_path)
+        end
+
+        model
+      end
+
+      # Load from saved model
+      def self.load(path)
+        raise ModelNotFoundError.new(path) unless File.directory?(path)
+
+        config = Hub::ConfigLoader.from_pretrained(path)
+        pooling_mode = (config.config["pooling_mode"] || "mean").to_sym
+
+        model = new(config, pooling_mode: pooling_mode)
+
+        weights_path = File.join(path, "model.safetensors")
+        Hub::SafetensorsLoader.load_into_model(model, weights_path, strict: false)
+
+        model
+      end
+
+      def initialize(config, pooling_mode: :mean)
+        super(config)
+
+        raise ArgumentError, "Invalid pooling mode: #{pooling_mode}" unless POOLING_MODES.include?(pooling_mode)
+
+        @pooling_mode = pooling_mode
+        @encoder = BertEncoder.new(config)
+
+        # Optional: normalize embeddings
+        @normalize = true
+      end
+
+      def forward(input_ids, attention_mask: nil, token_type_ids: nil)
+        encoder_output = @encoder.call(
+          input_ids,
+          attention_mask: attention_mask,
+          token_type_ids: token_type_ids
+        )
+
+        # Pool based on strategy
+        embeddings = case @pooling_mode
+        when :cls
+          encoder_output[:pooler_output]
+        when :mean
+          mean_pooling(encoder_output[:last_hidden_state], attention_mask)
+        when :max
+          max_pooling(encoder_output[:last_hidden_state], attention_mask)
+        end
+
+        # L2 normalize
+        embeddings = Torch::NN::Functional.normalize(embeddings, p: 2, dim: 1) if @normalize
+
+        { embeddings: embeddings }
+      end
+
+      # Encode texts to embeddings
+      #
+      # @param input_ids [Torch::Tensor] Token IDs
+      # @param attention_mask [Torch::Tensor] Attention mask
+      # @return [Torch::Tensor] Embeddings
+      def encode(input_ids, attention_mask: nil, token_type_ids: nil)
+        eval
+        Torch.no_grad do
+          output = forward(input_ids, attention_mask: attention_mask, token_type_ids: token_type_ids)
+          output[:embeddings]
+        end
+      end
+
+      # Compute similarity between two sets of embeddings
+      #
+      # @param embeddings_a [Torch::Tensor] First embeddings
+      # @param embeddings_b [Torch::Tensor] Second embeddings
+      # @return [Torch::Tensor] Similarity scores
+      def similarity(embeddings_a, embeddings_b)
+        # Cosine similarity (embeddings should be normalized)
+        Torch.matmul(embeddings_a, embeddings_b.transpose(0, 1))
+      end
+
+      def save(path)
+        FileUtils.mkdir_p(path)
+
+        weights_path = File.join(path, "model.safetensors")
+        Safetensors::Torch.save_file(state_dict, weights_path)
+
+        save_config = @config.to_h.merge(
+          "model_type" => "sentence_transformer",
+          "pooling_mode" => @pooling_mode.to_s
+        )
+
+        config_path = File.join(path, "config.json")
+        File.write(config_path, JSON.pretty_generate(save_config))
+      end
+
+      private
+
+      def mean_pooling(hidden_states, attention_mask)
+        if attention_mask
+          mask = attention_mask.unsqueeze(-1).expand_as(hidden_states).float
+          sum_embeddings = (hidden_states * mask).sum(dim: 1)
+          sum_mask = mask.sum(dim: 1).clamp(min: 1e-9)
+          sum_embeddings / sum_mask
+        else
+          hidden_states.mean(dim: 1)
+        end
+      end
+
+      def max_pooling(hidden_states, attention_mask)
+        if attention_mask
+          mask = attention_mask.unsqueeze(-1).expand_as(hidden_states)
+          hidden_states = hidden_states.masked_fill(mask == 0, -1e9)
+        end
+        hidden_states.max(dim: 1).values
+      end
+
+      def self.load_pretrained_weights(model, weights_path)
+        tensors = Safetensors::Torch.load_file(weights_path)
+
+        mapped = {}
+        model_keys = model.state_dict.keys
+
+        tensors.each do |name, tensor|
+          # Try direct mapping first
+          mapped_name = map_sentence_transformer_weight_name(name)
+
+          if model_keys.include?(mapped_name)
+            mapped[mapped_name] = tensor
+          end
+        end
+
+        # Use no_grad and copy! since torch.rb doesn't support strict: false
+        Torch.no_grad do
+          state_dict = model.state_dict
+          mapped.each do |name, tensor|
+            state_dict[name].copy!(tensor) if state_dict.key?(name)
+          end
+        end
+      end
+
+      def self.map_sentence_transformer_weight_name(hf_name)
+        name = hf_name.dup
+
+        # sentence-transformers uses "0." prefix for the encoder module
+        name = name.sub(/^0\./, "encoder.")
+
+        # Then apply BERT mappings
+        name = name.sub("auto_model.", "encoder.")
+        name = name.sub("bert.embeddings.word_embeddings", "encoder.word_embeddings")
+        name = name.sub("bert.embeddings.position_embeddings", "encoder.position_embeddings")
+        name = name.sub("bert.embeddings.token_type_embeddings", "encoder.token_type_embeddings")
+        name = name.sub("bert.embeddings.LayerNorm", "encoder.embeddings_layer_norm")
+        name = name.sub("embeddings.word_embeddings", "encoder.word_embeddings")
+        name = name.sub("embeddings.position_embeddings", "encoder.position_embeddings")
+        name = name.sub("embeddings.token_type_embeddings", "encoder.token_type_embeddings")
+        name = name.sub("embeddings.LayerNorm", "encoder.embeddings_layer_norm")
+        name = name.gsub("bert.encoder.layer", "encoder.layers")
+        name = name.gsub("encoder.layer", "encoder.layers")
+        name = name.gsub(".attention.self.query", ".attention.query")
+        name = name.gsub(".attention.self.key", ".attention.key")
+        name = name.gsub(".attention.self.value", ".attention.value")
+        name = name.gsub(".attention.output.dense", ".attention.out")
+        name = name.gsub(".attention.output.LayerNorm", ".attention_layer_norm")
+        name = name.gsub(".intermediate.dense", ".intermediate")
+        name = name.gsub(".output.dense", ".output")
+        name = name.gsub(".output.LayerNorm", ".output_layer_norm")
+        name = name.sub("bert.pooler.dense", "encoder.pooler_dense")
+        name = name.sub("pooler.dense", "encoder.pooler_dense")
+
+        name
+      end
+    end
+  end
+end

data/lib/fine/models/siglip2_for_image_classification.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+module Fine
+  module Models
+    # SigLIP2 model with classification head for image classification
+    class SigLIP2ForImageClassification < Base
+      attr_reader :encoder, :classifier, :num_labels
+
+      # Load a pretrained model from Hugging Face Hub
+      #
+      # @param model_id [String] Hugging Face model ID (e.g., "google/siglip2-base-patch16-224")
+      # @param num_labels [Integer] Number of classification labels
+      # @param freeze_encoder [Boolean] Whether to freeze the vision encoder
+      # @param dropout [Float] Dropout rate for classifier
+      # @return [SigLIP2ForImageClassification]
+      def self.from_pretrained(model_id, num_labels:, freeze_encoder: false, dropout: 0.0)
+        # Download model files
+        downloader = Hub::ModelDownloader.new(model_id)
+        model_path = downloader.download
+
+        # Load config
+        config = Hub::ConfigLoader.from_pretrained(model_path)
+
+        # Create model
+        model = new(config, num_labels: num_labels, dropout: dropout)
+
+        # Load pretrained weights into encoder
+        weights_path = downloader.file_path("model.safetensors")
+        load_result = Hub::SafetensorsLoader.load_into_model(
+          model.encoder,
+          weights_path,
+          strict: false,
+          prefix: "vision_model"
+        )
+
+        # Log any issues
+        if load_result[:missing_keys].any?
+          warn "Missing keys in encoder: #{load_result[:missing_keys].first(5).join(', ')}..."
+        end
+
+        # Freeze encoder if requested
+        model.encoder.freeze! if freeze_encoder
+
+        model
+      end
+
+      # Load a fine-tuned model from disk
+      #
+      # @param path [String] Path to saved model directory
+      # @return [SigLIP2ForImageClassification]
+      def self.load(path)
+        raise ModelNotFoundError.new(path, "Model directory not found") unless File.directory?(path)
+
+        # Load config
+        config = Hub::ConfigLoader.from_pretrained(path)
+        num_labels = config.config["num_labels"] || config.config["id2label"]&.size
+
+        raise ConfigurationError, "Cannot determine num_labels from saved model" unless num_labels
+
+        # Create model
+        model = new(config, num_labels: num_labels)
+
+        # Load weights
+        weights_path = File.join(path, "model.safetensors")
+        Hub::SafetensorsLoader.load_into_model(model, weights_path, strict: false)
+
+        model
+      end
+
+      def initialize(config, num_labels:, dropout: 0.0)
+        super(config)
+
+        @num_labels = num_labels
+
+        # Vision encoder
+        @encoder = SigLIP2VisionEncoder.new(config)
+
+        # Classification head
+        @classifier = ClassificationHead.new(
+          config.hidden_size,
+          num_labels,
+          dropout: dropout
+        )
+      end
+
+      def forward(pixel_values, labels: nil)
+        # Get image features from encoder
+        features = @encoder.call(pixel_values)
+
+        # Classification logits
+        logits = @classifier.call(features)
+
+        # Compute loss if labels provided
+        if labels
+          loss = Torch::NN::Functional.cross_entropy(logits, labels)
+          { loss: loss, logits: logits }
+        else
+          { logits: logits }
+        end
+      end
+
+      # Predict class for an image tensor
+      #
+      # @param pixel_values [Torch::Tensor] Image tensor (batch, C, H, W)
+      # @return [Torch::Tensor] Predicted class indices
+      def predict(pixel_values)
+        eval
+        Torch.no_grad do
+          output = forward(pixel_values)
+          output[:logits].argmax(dim: 1)
+        end
+      end
+
+      # Get probabilities for each class
+      #
+      # @param pixel_values [Torch::Tensor] Image tensor (batch, C, H, W)
+      # @return [Torch::Tensor] Class probabilities
+      def predict_proba(pixel_values)
+        eval
+        Torch.no_grad do
+          output = forward(pixel_values)
+          Torch::NN::Functional.softmax(output[:logits], dim: 1)
+        end
+      end
+
+      # Save the model to disk
+      #
+      # @param path [String] Directory path to save to
+      # @param label_map [Hash, nil] Optional mapping of label names to IDs
+      def save(path, label_map: nil)
+        FileUtils.mkdir_p(path)
+
+        # Save weights
+        weights_path = File.join(path, "model.safetensors")
+        Safetensors::Torch.save_file(state_dict, weights_path)
+
+        # Build config with num_labels
+        save_config = @config.to_h.merge(
+          "num_labels" => @num_labels,
+          "model_type" => "siglip2_image_classification"
+        )
+
+        # Add label mapping if provided
+        if label_map
+          save_config["id2label"] = label_map.invert.sort.to_h
+          save_config["label2id"] = label_map
+        end
+
+        # Save config
+        config_path = File.join(path, "config.json")
+        File.write(config_path, JSON.pretty_generate(save_config))
+      end
+    end
+  end
+end

data/lib/fine/models/siglip2_vision_encoder.rb

@@ -0,0 +1,190 @@
+# frozen_string_literal: true
+
+module Fine
+  module Models
+    # SigLIP2 Vision Transformer Encoder
+    #
+    # Implements the vision encoder portion of SigLIP2 for image feature extraction.
+    # Architecture follows the standard ViT with patch embedding, transformer blocks,
+    # and pooling.
+    class SigLIP2VisionEncoder < Base
+      def initialize(config)
+        super(config)
+
+        @hidden_size = config.hidden_size
+        @num_layers = config.num_hidden_layers
+        @num_heads = config.num_attention_heads
+        @intermediate_size = config.intermediate_size
+        @image_size = config.image_size
+        @patch_size = config.patch_size
+        @num_channels = config.num_channels
+        @layer_norm_eps = config.layer_norm_eps
+
+        # Patch embedding
+        @patch_embed = PatchEmbedding.new(
+          image_size: @image_size,
+          patch_size: @patch_size,
+          in_channels: @num_channels,
+          embed_dim: @hidden_size
+        )
+
+        # Position embedding (learnable)
+        num_patches = (@image_size / @patch_size) ** 2
+        @pos_embed = Torch::NN::Parameter.new(
+          Torch.zeros(1, num_patches, @hidden_size)
+        )
+
+        # Transformer blocks
+        @blocks = Torch::NN::ModuleList.new(
+          @num_layers.times.map do
+            TransformerBlock.new(
+              hidden_size: @hidden_size,
+              num_heads: @num_heads,
+              intermediate_size: @intermediate_size,
+              layer_norm_eps: @layer_norm_eps
+            )
+          end
+        )
+
+        # Final layer norm
+        @norm = Torch::NN::LayerNorm.new(@hidden_size, eps: @layer_norm_eps)
+
+        # Initialize position embedding
+        init_pos_embed
+      end
+
+      def forward(pixel_values)
+        # pixel_values: (batch, channels, height, width)
+
+        # Patch embedding: (batch, num_patches, hidden_size)
+        x = @patch_embed.call(pixel_values)
+
+        # Add position embedding
+        x = x + @pos_embed
+
+        # Transformer blocks
+        @blocks.each do |block|
+          x = block.call(x)
+        end
+
+        # Final layer norm
+        x = @norm.call(x)
+
+        # Pool: take mean of all patch embeddings
+        x.mean(dim: 1)
+      end
+
+      private
+
+      def init_pos_embed
+        # Initialize with normal distribution (truncated normal not available in torch.rb)
+        # The values will be overwritten when loading pretrained weights
+        Torch::NN::Init.normal!(@pos_embed, mean: 0.0, std: 0.02)
+      end
+    end
+
+    # Patch embedding layer
+    class PatchEmbedding < Torch::NN::Module
+      def initialize(image_size:, patch_size:, in_channels:, embed_dim:)
+        super()
+
+        @image_size = image_size
+        @patch_size = patch_size
+        @num_patches = (image_size / patch_size) ** 2
+
+        # Use conv2d for efficient patch extraction
+        @proj = Torch::NN::Conv2d.new(
+          in_channels, embed_dim, patch_size,
+          stride: patch_size
+        )
+      end
+
+      def forward(x)
+        # x: (batch, channels, height, width)
+        # output: (batch, num_patches, embed_dim)
+
+        x = @proj.call(x) # (batch, embed_dim, h/patch, w/patch)
+        x = x.flatten(2)   # (batch, embed_dim, num_patches)
+        x.transpose(1, 2)  # (batch, num_patches, embed_dim)
+      end
+    end
+
+    # Transformer block with self-attention and MLP
+    class TransformerBlock < Torch::NN::Module
+      def initialize(hidden_size:, num_heads:, intermediate_size:, layer_norm_eps:)
+        super()
+
+        @norm1 = Torch::NN::LayerNorm.new(hidden_size, eps: layer_norm_eps)
+        @attn = Attention.new(hidden_size: hidden_size, num_heads: num_heads)
+        @norm2 = Torch::NN::LayerNorm.new(hidden_size, eps: layer_norm_eps)
+        @mlp = MLP.new(hidden_size: hidden_size, intermediate_size: intermediate_size)
+      end
+
+      def forward(x)
+        # Pre-norm architecture
+        x = x + @attn.call(@norm1.call(x))
+        x = x + @mlp.call(@norm2.call(x))
+        x
+      end
+    end
+
+    # Multi-head self-attention
+    class Attention < Torch::NN::Module
+      def initialize(hidden_size:, num_heads:)
+        super()
+
+        @num_heads = num_heads
+        @head_dim = hidden_size / num_heads
+        @scale = @head_dim ** -0.5
+        @hidden_size = hidden_size
+
+        # Separate Q, K, V projections (matches HuggingFace)
+        @q_proj = Torch::NN::Linear.new(hidden_size, hidden_size)
+        @k_proj = Torch::NN::Linear.new(hidden_size, hidden_size)
+        @v_proj = Torch::NN::Linear.new(hidden_size, hidden_size)
+        @out_proj = Torch::NN::Linear.new(hidden_size, hidden_size)
+      end
+
+      def forward(x)
+        batch_size, seq_len, _ = x.shape
+
+        # Compute Q, K, V separately
+        q = @q_proj.call(x)
+        k = @k_proj.call(x)
+        v = @v_proj.call(x)
+
+        # Reshape for multi-head attention
+        q = q.reshape(batch_size, seq_len, @num_heads, @head_dim).transpose(1, 2)
+        k = k.reshape(batch_size, seq_len, @num_heads, @head_dim).transpose(1, 2)
+        v = v.reshape(batch_size, seq_len, @num_heads, @head_dim).transpose(1, 2)
+
+        # Scaled dot-product attention
+        attn = Torch.matmul(q, k.transpose(-2, -1)) * @scale
+        attn = attn.softmax(dim: -1)
+
+        # Apply attention to values
+        out = Torch.matmul(attn, v) # (batch, heads, seq, head_dim)
+        out = out.transpose(1, 2).reshape(batch_size, seq_len, @hidden_size)
+
+        @out_proj.call(out)
+      end
+    end
+
+    # MLP (feed-forward network)
+    class MLP < Torch::NN::Module
+      def initialize(hidden_size:, intermediate_size:)
+        super()
+
+        @fc1 = Torch::NN::Linear.new(hidden_size, intermediate_size)
+        @act = Torch::NN::GELU.new
+        @fc2 = Torch::NN::Linear.new(intermediate_size, hidden_size)
+      end
+
+      def forward(x)
+        x = @fc1.call(x)
+        x = @act.call(x)
+        @fc2.call(x)
+      end
+    end
+  end
+end