fine 0.1.0

data/lib/fine/text_classifier.rb

@@ -0,0 +1,250 @@
+# frozen_string_literal: true
+
+module Fine
+  # High-level API for text classification
+  #
+  # @example Basic usage
+  #   classifier = Fine::TextClassifier.new("distilbert-base-uncased")
+  #   classifier.fit(train_file: "reviews.jsonl", epochs: 3)
+  #   classifier.predict("This product is amazing!")
+  #
+  # @example With configuration
+  #   classifier = Fine::TextClassifier.new("microsoft/deberta-v3-small") do |config|
+  #     config.epochs = 5
+  #     config.batch_size = 16
+  #     config.learning_rate = 2e-5
+  #   end
+  #
+  class TextClassifier
+    attr_reader :model, :config, :tokenizer, :label_map, :model_id
+
+    # Create a new TextClassifier
+    #
+    # @param model_id [String] HuggingFace model ID
+    # @yield [config] Optional configuration block
+    def initialize(model_id, &block)
+      @model_id = model_id
+      @config = TextConfiguration.new
+      @model = nil
+      @tokenizer = nil
+      @label_map = nil
+      @trained = false
+
+      block&.call(@config)
+
+      if @config.callbacks.empty? && Fine.configuration&.progress_bar != false
+        @config.callbacks << Callbacks::ProgressBar.new
+      end
+    end
+
+    # Load a fine-tuned classifier from disk
+    #
+    # @param path [String] Path to saved model directory
+    # @return [TextClassifier]
+    def self.load(path)
+      config_path = File.join(path, "config.json")
+      raise ModelNotFoundError.new(path) unless File.exist?(config_path)
+
+      config_data = JSON.parse(File.read(config_path))
+
+      classifier = allocate
+      classifier.instance_variable_set(:@model_id, config_data["_model_id"] || "custom")
+      classifier.instance_variable_set(:@config, TextConfiguration.new)
+      classifier.instance_variable_set(:@trained, true)
+
+      # Load label map
+      if config_data["label2id"]
+        classifier.instance_variable_set(:@label_map, config_data["label2id"])
+      elsif config_data["id2label"]
+        label_map = config_data["id2label"].transform_keys(&:to_i).invert
+        classifier.instance_variable_set(:@label_map, label_map)
+      end
+
+      # Load tokenizer
+      tokenizer_path = File.join(path, "tokenizer.json")
+      tokenizer = if File.exist?(tokenizer_path)
+        Tokenizers::AutoTokenizer.new(path, max_length: config_data["max_length"] || 512)
+      else
+        Tokenizers::AutoTokenizer.from_pretrained(
+          config_data["_model_id"] || "distilbert-base-uncased",
+          max_length: config_data["max_length"] || 512
+        )
+      end
+      classifier.instance_variable_set(:@tokenizer, tokenizer)
+
+      # Load model
+      classifier.instance_variable_set(
+        :@model,
+        Models::BertForSequenceClassification.load(path)
+      )
+
+      classifier
+    end
+
+    # Fine-tune on a dataset
+    #
+    # @param train_file [String] Path to training data (JSONL or CSV)
+    # @param val_file [String, nil] Path to validation data
+    # @param epochs [Integer, nil] Override config epochs
+    # @return [Array<Hash>] Training history
+    def fit(train_file:, val_file: nil, epochs: nil)
+      @config.epochs = epochs if epochs
+
+      # Load tokenizer
+      @tokenizer = Tokenizers::AutoTokenizer.from_pretrained(
+        @model_id,
+        max_length: @config.max_length
+      )
+
+      # Load datasets
+      train_dataset = Datasets::TextDataset.from_file(train_file, tokenizer: @tokenizer)
+
+      val_dataset = if val_file
+        Datasets::TextDataset.from_file(val_file, tokenizer: @tokenizer, label_map: train_dataset.label_map)
+      end
+
+      @label_map = train_dataset.label_map
+      num_classes = train_dataset.num_classes
+
+      # Load model
+      @model = Models::BertForSequenceClassification.from_pretrained(
+        @model_id,
+        num_labels: num_classes,
+        dropout: @config.dropout
+      )
+
+      # Train
+      trainer = Training::TextTrainer.new(
+        @model,
+        @config,
+        train_dataset: train_dataset,
+        val_dataset: val_dataset
+      )
+
+      history = trainer.fit
+      @trained = true
+
+      history
+    end
+
+    # Fine-tune with automatic train/val split
+    #
+    # @param data_file [String] Path to data file
+    # @param val_split [Float] Fraction for validation
+    # @return [Array<Hash>] Training history
+    def fit_with_split(data_file:, val_split: 0.2, epochs: nil)
+      @config.epochs = epochs if epochs
+
+      @tokenizer = Tokenizers::AutoTokenizer.from_pretrained(
+        @model_id,
+        max_length: @config.max_length
+      )
+
+      full_dataset = Datasets::TextDataset.from_file(data_file, tokenizer: @tokenizer)
+      train_dataset, val_dataset = full_dataset.split(test_size: val_split)
+
+      @label_map = train_dataset.label_map
+
+      @model = Models::BertForSequenceClassification.from_pretrained(
+        @model_id,
+        num_labels: train_dataset.num_classes,
+        dropout: @config.dropout
+      )
+
+      trainer = Training::TextTrainer.new(
+        @model,
+        @config,
+        train_dataset: train_dataset,
+        val_dataset: val_dataset
+      )
+
+      history = trainer.fit
+      @trained = true
+
+      history
+    end
+
+    # Make predictions
+    #
+    # @param texts [String, Array<String>] Text(s) to classify
+    # @param top_k [Integer] Number of top predictions to return
+    # @return [Array<Array<Hash>>] Predictions with :label and :score
+    def predict(texts, top_k: 5)
+      raise TrainingError, "Model not trained or loaded" unless @trained && @model
+
+      texts = [texts] if texts.is_a?(String)
+
+      # Tokenize
+      encoding = @tokenizer.encode(texts)
+
+      # Get predictions
+      @model.eval
+      probs = @model.predict_proba(
+        encoding[:input_ids],
+        attention_mask: encoding[:attention_mask],
+        token_type_ids: encoding[:token_type_ids]
+      )
+
+      # Convert to result format
+      inverse_label_map = @label_map.invert
+
+      probs.to_a.map do |sample_probs|
+        sorted = sample_probs.each_with_index.sort_by { |prob, _| -prob }
+        top = sorted.first([top_k, @label_map.size].min)
+
+        top.map do |prob, idx|
+          {
+            label: inverse_label_map[idx] || idx.to_s,
+            score: prob.round(4)
+          }
+        end
+      end
+    end
+
+    # Save the model
+    #
+    # @param path [String] Directory to save to
+    def save(path)
+      raise TrainingError, "Model not trained" unless @trained && @model
+
+      @model.save(path, label_map: @label_map)
+      @tokenizer.save(path)
+
+      # Update config with model ID and max_length
+      config_path = File.join(path, "config.json")
+      config = JSON.parse(File.read(config_path))
+      config["_model_id"] = @model_id
+      config["max_length"] = @config.max_length
+      File.write(config_path, JSON.pretty_generate(config))
+    end
+
+    # Get class names
+    def class_names
+      return [] unless @label_map
+
+      @label_map.sort_by { |_, v| v }.map(&:first)
+    end
+
+    # Export to ONNX format
+    #
+    # @param path [String] Output path for ONNX file
+    # @param options [Hash] Export options
+    # @return [String] The output path
+    def export_onnx(path, **options)
+      Export.to_onnx(self, path, **options)
+    end
+  end
+
+  # Configuration for text models
+  class TextConfiguration < Configuration
+    attr_accessor :max_length, :warmup_ratio
+
+    def initialize
+      super
+      @max_length = 256
+      @warmup_ratio = 0.1
+      @learning_rate = 2e-5  # Lower default for text models
+      @batch_size = 16       # Smaller default for text
+    end
+  end
+end

data/lib/fine/text_embedder.rb

@@ -0,0 +1,221 @@
+# frozen_string_literal: true
+
+module Fine
+  # High-level API for text embeddings
+  #
+  # @example Basic usage
+  #   embedder = Fine::TextEmbedder.new("sentence-transformers/all-MiniLM-L6-v2")
+  #   embedder.fit(train_file: "pairs.jsonl", epochs: 3)
+  #   embedding = embedder.encode("Hello world")
+  #
+  # @example Without fine-tuning (use pretrained directly)
+  #   embedder = Fine::TextEmbedder.new("sentence-transformers/all-MiniLM-L6-v2")
+  #   embedding = embedder.encode("Hello world")
+  #
+  class TextEmbedder
+    attr_reader :model, :config, :tokenizer, :model_id
+
+    # Create a new TextEmbedder
+    #
+    # @param model_id [String] HuggingFace model ID
+    # @yield [config] Optional configuration block
+    def initialize(model_id, &block)
+      @model_id = model_id
+      @config = EmbeddingConfiguration.new
+      @model = nil
+      @tokenizer = nil
+      @trained = false
+
+      block&.call(@config)
+
+      # Load tokenizer immediately for encoding
+      @tokenizer = Tokenizers::AutoTokenizer.from_pretrained(
+        model_id,
+        max_length: @config.max_length
+      )
+
+      # Load pretrained model for immediate use
+      @model = Models::SentenceTransformer.from_pretrained(
+        model_id,
+        pooling_mode: @config.pooling_mode
+      )
+      @trained = true  # Pretrained is ready to use
+    end
+
+    # Load a fine-tuned embedder from disk
+    #
+    # @param path [String] Path to saved model
+    # @return [TextEmbedder]
+    def self.load(path)
+      config_path = File.join(path, "config.json")
+      raise ModelNotFoundError.new(path) unless File.exist?(config_path)
+
+      config_data = JSON.parse(File.read(config_path))
+
+      embedder = allocate
+      embedder.instance_variable_set(:@model_id, config_data["_model_id"] || "custom")
+      embedder.instance_variable_set(:@config, EmbeddingConfiguration.new)
+      embedder.instance_variable_set(:@trained, true)
+
+      # Load tokenizer
+      tokenizer_path = File.join(path, "tokenizer.json")
+      tokenizer = if File.exist?(tokenizer_path)
+        Tokenizers::AutoTokenizer.new(path, max_length: config_data["max_length"] || 512)
+      else
+        Tokenizers::AutoTokenizer.from_pretrained(
+          config_data["_model_id"],
+          max_length: config_data["max_length"] || 512
+        )
+      end
+      embedder.instance_variable_set(:@tokenizer, tokenizer)
+
+      # Load model
+      embedder.instance_variable_set(
+        :@model,
+        Models::SentenceTransformer.load(path)
+      )
+
+      embedder
+    end
+
+    # Fine-tune on pairs/triplets data
+    #
+    # @param train_file [String] Path to training data (JSONL)
+    # @param epochs [Integer, nil] Override config epochs
+    # @return [Array<Hash>] Training history
+    def fit(train_file:, epochs: nil)
+      @config.epochs = epochs if epochs
+
+      # Load dataset
+      train_dataset = Datasets::TextPairDataset.from_jsonl(
+        train_file,
+        tokenizer: @tokenizer,
+        text_a_column: "query",
+        text_b_column: "positive"
+      )
+
+      # Add progress bar callback
+      if @config.callbacks.empty? && Fine.configuration&.progress_bar != false
+        @config.callbacks << Callbacks::ProgressBar.new
+      end
+
+      # Train
+      trainer = Training::EmbeddingTrainer.new(
+        @model,
+        @config,
+        train_dataset: train_dataset
+      )
+
+      trainer.fit
+    end
+
+    # Encode text(s) to embeddings
+    #
+    # @param texts [String, Array<String>] Text(s) to encode
+    # @return [Array<Float>, Array<Array<Float>>] Embedding(s)
+    def encode(texts)
+      raise TrainingError, "Model not loaded" unless @model
+
+      single_input = texts.is_a?(String)
+      texts = [texts] if single_input
+
+      # Tokenize
+      encoding = @tokenizer.encode(texts)
+
+      # Get embeddings
+      embeddings = @model.encode(
+        encoding[:input_ids],
+        attention_mask: encoding[:attention_mask]
+      )
+
+      # Convert to Ruby arrays
+      result = embeddings.to_a
+
+      single_input ? result.first : result
+    end
+
+    # Compute similarity between two texts
+    #
+    # @param text_a [String] First text
+    # @param text_b [String] Second text
+    # @return [Float] Cosine similarity score
+    def similarity(text_a, text_b)
+      emb_a = encode(text_a)
+      emb_b = encode(text_b)
+
+      cosine_similarity(emb_a, emb_b)
+    end
+
+    # Find most similar texts from a corpus
+    #
+    # @param query [String] Query text
+    # @param corpus [Array<String>] Corpus to search
+    # @param top_k [Integer] Number of results
+    # @return [Array<Hash>] Results with :text, :score, :index
+    def search(query, corpus, top_k: 5)
+      query_emb = encode(query)
+      corpus_embs = encode(corpus)
+
+      scores = corpus_embs.map.with_index do |emb, idx|
+        { text: corpus[idx], score: cosine_similarity(query_emb, emb), index: idx }
+      end
+
+      scores.sort_by { |s| -s[:score] }.first(top_k)
+    end
+
+    # Save the model
+    #
+    # @param path [String] Directory to save to
+    def save(path)
+      raise TrainingError, "Model not loaded" unless @model
+
+      @model.save(path)
+      @tokenizer.save(path)
+
+      # Update config
+      config_path = File.join(path, "config.json")
+      config = JSON.parse(File.read(config_path))
+      config["_model_id"] = @model_id
+      config["max_length"] = @config.max_length
+      File.write(config_path, JSON.pretty_generate(config))
+    end
+
+    # Get embedding dimension
+    def embedding_dim
+      @model.config.hidden_size
+    end
+
+    # Export to ONNX format
+    #
+    # @param path [String] Output path for ONNX file
+    # @param options [Hash] Export options
+    # @return [String] The output path
+    def export_onnx(path, **options)
+      Export.to_onnx(self, path, **options)
+    end
+
+    private
+
+    def cosine_similarity(a, b)
+      dot = a.zip(b).sum { |x, y| x * y }
+      norm_a = Math.sqrt(a.sum { |x| x * x })
+      norm_b = Math.sqrt(b.sum { |x| x * x })
+      dot / (norm_a * norm_b)
+    end
+  end
+
+  # Configuration for embedding models
+  class EmbeddingConfiguration < Configuration
+    attr_accessor :max_length, :pooling_mode, :loss
+
+    def initialize
+      super
+      @max_length = 256
+      @pooling_mode = :mean
+      @loss = :multiple_negatives_ranking
+      @learning_rate = 2e-5
+      @batch_size = 32
+      @epochs = 1
+    end
+  end
+end

data/lib/fine/tokenizers/auto_tokenizer.rb

@@ -0,0 +1,208 @@
+# frozen_string_literal: true
+
+require "tokenizers"
+
+module Fine
+  module Tokenizers
+    # Wrapper around HuggingFace tokenizers
+    class AutoTokenizer
+      attr_reader :tokenizer, :model_id, :max_length
+
+      # Load tokenizer from a pretrained model
+      #
+      # @param model_id [String] HuggingFace model ID
+      # @param max_length [Integer] Maximum sequence length
+      # @return [AutoTokenizer]
+      def self.from_pretrained(model_id, max_length: 512)
+        new(model_id, max_length: max_length)
+      end
+
+      def initialize(model_id, max_length: 512)
+        @model_id = model_id
+        @max_length = max_length
+        @tokenizer = load_tokenizer(model_id)
+
+        configure_tokenizer
+      end
+
+      # Tokenize a single text or batch of texts
+      #
+      # @param texts [String, Array<String>] Text(s) to tokenize
+      # @param padding [Boolean] Whether to pad sequences
+      # @param truncation [Boolean] Whether to truncate sequences
+      # @param return_tensors [Boolean] Whether to return Torch tensors
+      # @return [Hash] Hash with :input_ids, :attention_mask, and optionally :token_type_ids
+      def encode(texts, padding: true, truncation: true, return_tensors: true)
+        texts = [texts] if texts.is_a?(String)
+        single_input = texts.size == 1
+
+        # Encode all texts
+        encodings = texts.map do |text|
+          @tokenizer.encode(text)
+        end
+
+        # Get max length in batch for padding
+        max_len = if padding
+          [encodings.map { |e| e.ids.length }.max, @max_length].min
+        else
+          @max_length
+        end
+
+        # Build output arrays
+        input_ids = []
+        attention_mask = []
+        token_type_ids = []
+
+        encodings.each do |encoding|
+          ids = encoding.ids
+          mask = encoding.attention_mask
+          type_ids = encoding.type_ids rescue Array.new(ids.length, 0)
+
+          # Truncate if needed
+          if truncation && ids.length > max_len
+            ids = ids[0...max_len]
+            mask = mask[0...max_len]
+            type_ids = type_ids[0...max_len]
+          end
+
+          # Pad if needed
+          if padding && ids.length < max_len
+            pad_length = max_len - ids.length
+            ids = ids + Array.new(pad_length, pad_token_id)
+            mask = mask + Array.new(pad_length, 0)
+            type_ids = type_ids + Array.new(pad_length, 0)
+          end
+
+          input_ids << ids
+          attention_mask << mask
+          token_type_ids << type_ids
+        end
+
+        result = {
+          input_ids: input_ids,
+          attention_mask: attention_mask
+        }
+
+        # Only include token_type_ids for BERT-style models
+        result[:token_type_ids] = token_type_ids if has_token_type_ids?
+
+        if return_tensors
+          result.transform_values! { |v| Torch.tensor(v, dtype: :long) }
+        end
+
+        result
+      end
+
+      # Encode a pair of texts (for sentence pair tasks)
+      #
+      # @param text_a [String] First text
+      # @param text_b [String] Second text
+      # @return [Hash] Tokenized output
+      def encode_pair(text_a, text_b, **kwargs)
+        encoding = @tokenizer.encode(text_a, text_b)
+
+        ids = encoding.ids
+        mask = encoding.attention_mask
+        type_ids = encoding.type_ids rescue Array.new(ids.length, 0)
+
+        # Truncate if needed
+        if kwargs.fetch(:truncation, true) && ids.length > @max_length
+          ids = ids[0...@max_length]
+          mask = mask[0...@max_length]
+          type_ids = type_ids[0...@max_length]
+        end
+
+        result = {
+          input_ids: [ids],
+          attention_mask: [mask]
+        }
+        result[:token_type_ids] = [type_ids] if has_token_type_ids?
+
+        if kwargs.fetch(:return_tensors, true)
+          result.transform_values! { |v| Torch.tensor(v, dtype: :long) }
+        end
+
+        result
+      end
+
+      # Decode token IDs back to text
+      #
+      # @param token_ids [Array<Integer>] Token IDs
+      # @param skip_special_tokens [Boolean] Whether to skip special tokens
+      # @return [String] Decoded text
+      def decode(token_ids, skip_special_tokens: true)
+        token_ids = token_ids.to_a if token_ids.respond_to?(:to_a)
+        @tokenizer.decode(token_ids, skip_special_tokens: skip_special_tokens)
+      end
+
+      # Get vocabulary size
+      def vocab_size
+        @tokenizer.vocab_size
+      end
+
+      # Get pad token ID
+      def pad_token_id
+        @tokenizer.token_to_id(@tokenizer.padding&.dig("pad_token") || "[PAD]") || 0
+      end
+
+      # Get CLS token ID
+      def cls_token_id
+        @tokenizer.token_to_id("[CLS]") || @tokenizer.token_to_id("<s>") || 0
+      end
+
+      # Get SEP token ID
+      def sep_token_id
+        @tokenizer.token_to_id("[SEP]") || @tokenizer.token_to_id("</s>") || 0
+      end
+
+      # Get EOS token ID
+      def eos_token_id
+        @tokenizer.token_to_id("</s>") || @tokenizer.token_to_id("[SEP]") || @tokenizer.token_to_id("<|endoftext|>") || 0
+      end
+
+      # Get BOS token ID
+      def bos_token_id
+        @tokenizer.token_to_id("<s>") || @tokenizer.token_to_id("[CLS]") || @tokenizer.token_to_id("<|startoftext|>") || 0
+      end
+
+      # Save tokenizer to directory
+      def save(path)
+        FileUtils.mkdir_p(path)
+        @tokenizer.save(File.join(path, "tokenizer.json"))
+      end
+
+      private
+
+      def load_tokenizer(model_id)
+        # Check if it's a local path with tokenizer.json
+        local_tokenizer_path = File.join(model_id, "tokenizer.json")
+        if File.exist?(local_tokenizer_path)
+          return ::Tokenizers::Tokenizer.from_file(local_tokenizer_path)
+        end
+
+        # Check if model_id itself is a tokenizer.json path
+        if File.exist?(model_id) && model_id.end_with?("tokenizer.json")
+          return ::Tokenizers::Tokenizer.from_file(model_id)
+        end
+
+        # Try to load from HuggingFace Hub
+        ::Tokenizers::Tokenizer.from_pretrained(model_id)
+      rescue StandardError => e
+        raise ConfigurationError, "Failed to load tokenizer for #{model_id}: #{e.message}"
+      end
+
+      def configure_tokenizer
+        # Enable truncation
+        @tokenizer.enable_truncation(@max_length)
+
+        # Enable padding
+        @tokenizer.enable_padding(length: @max_length)
+      end
+
+      def has_token_type_ids?
+        # BERT-style models use token_type_ids, RoBERTa/DistilBERT don't always need them
+        @model_id.include?("bert") && !@model_id.include?("roberta")
+      end
+    end
+  end
+end