fastembed 1.0.0 → 1.1.0

data/lib/fastembed/onnx_embedding_model.rb

@@ -5,20 +5,41 @@ require 'tokenizers'
 
 module Fastembed
   # ONNX-based embedding model wrapper
+  #
+  # Handles the low-level details of running ONNX inference and tokenization
+  # for embedding models. Used internally by TextEmbedding.
+  #
+  # @api private
+  #
   class OnnxEmbeddingModel
+    # @!attribute [r] model_info
+    #   @return [ModelInfo] Model metadata
+    # @!attribute [r] model_dir
+    #   @return [String] Path to model directory
     attr_reader :model_info, :model_dir
 
-    def initialize(model_info, model_dir, threads: nil, providers: nil)
+    # @param model_info [ModelInfo] Model metadata
+    # @param model_dir [String] Directory containing model files
+    # @param threads [Integer, nil] Number of threads for ONNX Runtime
+    # @param providers [Array<String>, nil] ONNX execution providers
+    # @param model_file_override [String, nil] Override model file path (for quantized models)
+    def initialize(model_info, model_dir, threads: nil, providers: nil, model_file_override: nil)
       @model_info = model_info
       @model_dir = model_dir
       @threads = threads
       @providers = providers
+      @model_file = model_file_override || model_info.model_file
 
       load_model
       load_tokenizer
     end
 
     # Embed a batch of texts
+    #
+    # Tokenizes the input, runs ONNX inference, and applies pooling.
+    #
+    # @param texts [Array<String>] Texts to embed
+    # @return [Array<Array<Float>>] Embedding vectors
     def embed(texts)
       # Tokenize
       encoded = tokenize(texts)
@@ -38,7 +59,7 @@ module Fastembed
     private
 
     def load_model
-      model_path = File.join(model_dir, model_info.model_file)
+      model_path = File.join(model_dir, @model_file)
       raise Error, "Model file not found: #{model_path}" unless File.exist?(model_path)
 
       session_options = {}
@@ -53,11 +74,11 @@ module Fastembed
       tokenizer_path = File.join(model_dir, model_info.tokenizer_file)
       raise Error, "Tokenizer file not found: #{tokenizer_path}" unless File.exist?(tokenizer_path)
 
-      @tokenizer = Tokenizers.from_file(tokenizer_path)
+      @tokenizer = Tokenizers::Tokenizer.from_file(tokenizer_path)
 
       # Configure tokenizer for batch encoding
       @tokenizer.enable_padding(pad_id: 0, pad_token: '[PAD]')
-      @tokenizer.enable_truncation(512)
+      @tokenizer.enable_truncation(model_info.max_length)
     end
 
     def tokenize(texts)

data/lib/fastembed/pooling.rb

@@ -1,10 +1,32 @@
 # frozen_string_literal: true
 
 module Fastembed
-  # Pooling strategies for transformer outputs
+  # Pooling strategies for transformer model outputs
+  #
+  # Transforms variable-length token embeddings into fixed-size sentence embeddings.
+  # Supports mean pooling (default) and CLS token pooling.
+  #
+  # @example Apply mean pooling with normalization
+  #   pooled = Pooling.apply(:mean, token_embeddings, attention_mask)
+  #
   module Pooling
+    # Valid pooling strategies
+    VALID_STRATEGIES = %i[mean cls].freeze
+
+    # Check if a pooling strategy is valid
+    #
+    # @param strategy [Symbol] Pooling strategy to check
+    # @return [Boolean] True if valid
+    def self.valid?(strategy)
+      VALID_STRATEGIES.include?(strategy)
+    end
+
     class << self
       # Mean pooling - averages all token embeddings weighted by attention mask
+      #
+      # @param token_embeddings [Array<Array<Array<Float>>>] Token embeddings [batch, seq, hidden]
+      # @param attention_mask [Array<Array<Integer>>] Attention mask [batch, seq]
+      # @return [Array<Array<Float>>] Pooled embeddings [batch, hidden]
       def mean_pooling(token_embeddings, attention_mask)
         # token_embeddings: [batch_size, seq_len, hidden_size]
         # attention_mask: [batch_size, seq_len]
@@ -40,11 +62,18 @@ module Fastembed
       end
 
       # CLS pooling - uses the [CLS] token embedding (first token)
+      #
+      # @param token_embeddings [Array<Array<Array<Float>>>] Token embeddings [batch, seq, hidden]
+      # @param _attention_mask [Array<Array<Integer>>] Attention mask (unused)
+      # @return [Array<Array<Float>>] Pooled embeddings [batch, hidden]
       def cls_pooling(token_embeddings, _attention_mask)
         token_embeddings.map { |batch| batch[0] }
       end
 
-      # L2 normalize vectors
+      # L2 normalize vectors to unit length
+      #
+      # @param vectors [Array<Array<Float>>] Vectors to normalize
+      # @return [Array<Array<Float>>] Normalized vectors
       def normalize(vectors)
         vectors.map do |vector|
           norm = Math.sqrt(vector.sum { |v| v * v })
@@ -53,7 +82,14 @@ module Fastembed
         end
       end
 
-      # Apply pooling based on strategy
+      # Apply pooling strategy to token embeddings
+      #
+      # @param strategy [Symbol] Pooling strategy (:mean or :cls)
+      # @param token_embeddings [Array] Token embeddings from model
+      # @param attention_mask [Array] Attention mask
+      # @param should_normalize [Boolean] Whether to L2 normalize output
+      # @return [Array<Array<Float>>] Pooled embeddings
+      # @raise [ArgumentError] If unknown pooling strategy
       def apply(strategy, token_embeddings, attention_mask, should_normalize: true)
         pooled = case strategy
                  when :mean

data/lib/fastembed/progress.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Fastembed
+  # Progress information for batch operations
+  # Passed to progress callbacks during embedding
+  class Progress
+    attr_reader :current, :total, :batch_size
+
+    # @param current [Integer] Current batch number (1-indexed)
+    # @param total [Integer] Total number of batches
+    # @param batch_size [Integer] Size of each batch
+    def initialize(current:, total:, batch_size:)
+      @current = current
+      @total = total
+      @batch_size = batch_size
+    end
+
+    # Percentage complete (0.0 to 1.0)
+    # @return [Float]
+    def percentage
+      return 1.0 if total.zero?
+
+      current.to_f / total
+    end
+
+    # Percentage as integer (0 to 100)
+    # @return [Integer]
+    def percent
+      (percentage * 100).round
+    end
+
+    # Number of documents processed so far
+    # @return [Integer]
+    def documents_processed
+      current * batch_size
+    end
+
+    # Check if processing is complete
+    # @return [Boolean]
+    def complete?
+      current >= total
+    end
+
+    def to_s
+      "Progress(#{current}/#{total}, #{percent}%)"
+    end
+
+    def inspect
+      to_s
+    end
+  end
+end

data/lib/fastembed/quantization.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module Fastembed
+  # Quantization options for ONNX models
+  # Different quantization levels trade off model size/speed vs accuracy
+  module Quantization
+    # Available quantization types
+    TYPES = {
+      fp32: {
+        suffix: '',
+        description: 'Full precision (32-bit float)',
+        size_multiplier: 1.0
+      },
+      fp16: {
+        suffix: '_fp16',
+        description: 'Half precision (16-bit float)',
+        size_multiplier: 0.5
+      },
+      int8: {
+        suffix: '_int8',
+        description: 'Dynamic int8 quantization',
+        size_multiplier: 0.25
+      },
+      uint8: {
+        suffix: '_uint8',
+        description: 'Dynamic uint8 quantization',
+        size_multiplier: 0.25
+      },
+      q4: {
+        suffix: '_q4',
+        description: '4-bit quantization',
+        size_multiplier: 0.125
+      }
+    }.freeze
+
+    # Default quantization type
+    DEFAULT = :fp32
+
+    class << self
+      # Check if a quantization type is valid
+      # @param type [Symbol] Quantization type
+      # @return [Boolean]
+      def valid?(type)
+        TYPES.key?(type)
+      end
+
+      # Get the model file suffix for a quantization type
+      # @param type [Symbol] Quantization type
+      # @return [String] Suffix to append to model filename
+      def suffix(type)
+        TYPES.dig(type, :suffix) || ''
+      end
+
+      # Get quantized model filename
+      # @param base_file [String] Base model file (e.g., "onnx/model.onnx")
+      # @param type [Symbol] Quantization type
+      # @return [String] Quantized model filename
+      def model_file(base_file, type)
+        return base_file if type == :fp32 || type.nil?
+
+        ext = File.extname(base_file)
+        base = base_file.chomp(ext)
+        "#{base}#{suffix(type)}#{ext}"
+      end
+
+      # List available quantization types
+      # @return [Array<Hash>] Array of quantization info
+      def list
+        TYPES.map do |type, info|
+          { type: type, description: info[:description], size_multiplier: info[:size_multiplier] }
+        end
+      end
+    end
+  end
+end

data/lib/fastembed/reranker_model_info.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module Fastembed
+  # Model information for reranker/cross-encoder models
+  #
+  # Cross-encoders process query-document pairs together rather than
+  # encoding them separately, enabling more accurate relevance scoring.
+  #
+  # @example Access reranker info
+  #   info = Fastembed::SUPPORTED_RERANKER_MODELS['BAAI/bge-reranker-base']
+  #   info.model_name  # => "BAAI/bge-reranker-base"
+  #
+  class RerankerModelInfo
+    include BaseModelInfo
+
+    # Create a new RerankerModelInfo instance
+    #
+    # @param model_name [String] Full model identifier
+    # @param description [String] Human-readable description
+    # @param size_in_gb [Float] Model size in GB
+    # @param sources [Hash] Source repositories
+    # @param model_file [String] Path to ONNX model file
+    # @param tokenizer_file [String] Path to tokenizer file
+    def initialize(
+      model_name:,
+      description:,
+      size_in_gb:,
+      sources:,
+      model_file: 'onnx/model.onnx',
+      tokenizer_file: 'tokenizer.json'
+    )
+      initialize_base(
+        model_name: model_name,
+        description: description,
+        size_in_gb: size_in_gb,
+        sources: sources,
+        model_file: model_file,
+        tokenizer_file: tokenizer_file
+      )
+    end
+
+    # Convert to hash representation
+    # @return [Hash] Model info as a hash
+    def to_h
+      {
+        model_name: model_name,
+        description: description,
+        size_in_gb: size_in_gb,
+        sources: sources,
+        model_file: model_file,
+        tokenizer_file: tokenizer_file
+      }
+    end
+  end
+
+  # Registry of supported reranker models
+  SUPPORTED_RERANKER_MODELS = {
+    'BAAI/bge-reranker-base' => RerankerModelInfo.new(
+      model_name: 'BAAI/bge-reranker-base',
+      description: 'BGE reranker base model for relevance scoring',
+      size_in_gb: 1.11,
+      sources: { hf: 'Xenova/bge-reranker-base' }
+    ),
+    'BAAI/bge-reranker-large' => RerankerModelInfo.new(
+      model_name: 'BAAI/bge-reranker-large',
+      description: 'BGE reranker large model for higher accuracy',
+      size_in_gb: 2.24,
+      sources: { hf: 'Xenova/bge-reranker-large' }
+    ),
+    'cross-encoder/ms-marco-MiniLM-L-6-v2' => RerankerModelInfo.new(
+      model_name: 'cross-encoder/ms-marco-MiniLM-L-6-v2',
+      description: 'Lightweight MS MARCO cross-encoder',
+      size_in_gb: 0.08,
+      sources: { hf: 'Xenova/ms-marco-MiniLM-L-6-v2' }
+    ),
+    'cross-encoder/ms-marco-MiniLM-L-12-v2' => RerankerModelInfo.new(
+      model_name: 'cross-encoder/ms-marco-MiniLM-L-12-v2',
+      description: 'MS MARCO cross-encoder with better accuracy',
+      size_in_gb: 0.12,
+      sources: { hf: 'Xenova/ms-marco-MiniLM-L-12-v2' }
+    ),
+    'jinaai/jina-reranker-v1-turbo-en' => RerankerModelInfo.new(
+      model_name: 'jinaai/jina-reranker-v1-turbo-en',
+      description: 'Fast Jina reranker for English',
+      size_in_gb: 0.13,
+      sources: { hf: 'jinaai/jina-reranker-v1-turbo-en' }
+    )
+  }.freeze
+
+  DEFAULT_RERANKER_MODEL = 'cross-encoder/ms-marco-MiniLM-L-6-v2'
+end

data/lib/fastembed/sparse_embedding.rb

@@ -0,0 +1,261 @@
+# frozen_string_literal: true
+
+module Fastembed
+  # Represents a sparse embedding vector
+  # Contains indices and their corresponding values (non-zero only)
+  class SparseEmbedding
+    attr_reader :indices, :values
+
+    def initialize(indices:, values:)
+      @indices = indices
+      @values = values
+    end
+
+    # Convert to hash format {index => value}
+    def to_h
+      indices.zip(values).to_h
+    end
+
+    # Number of non-zero elements
+    def nnz
+      indices.length
+    end
+
+    def to_s
+      "SparseEmbedding(nnz=#{nnz})"
+    end
+
+    def inspect
+      to_s
+    end
+  end
+
+  # Sparse text embedding using SPLADE models
+  #
+  # SPLADE (Sparse Lexical and Expansion) models produce sparse vectors
+  # where each dimension corresponds to a vocabulary token.
+  # These are useful for hybrid search combining with dense embeddings.
+  #
+  # @example Basic usage
+  #   sparse = Fastembed::TextSparseEmbedding.new
+  #   embeddings = sparse.embed(["Hello world"]).to_a
+  #   embeddings.first.indices  # => [101, 7592, 2088, ...]
+  #   embeddings.first.values   # => [0.45, 1.23, 0.89, ...]
+  #
+  class TextSparseEmbedding
+    include BaseModel
+
+    # Initialize a sparse embedding model
+    #
+    # @param model_name [String] Name of the model to use
+    # @param cache_dir [String, nil] Custom cache directory for models
+    # @param threads [Integer, nil] Number of threads for ONNX Runtime
+    # @param providers [Array<String>, nil] ONNX execution providers
+    # @param show_progress [Boolean] Whether to show download progress
+    # @param quantization [Symbol] Quantization type (:fp32, :fp16, :int8, :uint8, :q4)
+    # @param local_model_dir [String, nil] Load model from local directory instead of downloading
+    # @param model_file [String, nil] Override model file name (e.g., "model.onnx")
+    # @param tokenizer_file [String, nil] Override tokenizer file name (e.g., "tokenizer.json")
+    def initialize(
+      model_name: DEFAULT_SPARSE_MODEL,
+      cache_dir: nil,
+      threads: nil,
+      providers: nil,
+      show_progress: true,
+      quantization: nil,
+      local_model_dir: nil,
+      model_file: nil,
+      tokenizer_file: nil
+    )
+      if local_model_dir
+        initialize_from_local(
+          local_model_dir: local_model_dir,
+          model_name: model_name,
+          threads: threads,
+          providers: providers,
+          quantization: quantization,
+          model_file: model_file,
+          tokenizer_file: tokenizer_file
+        )
+      else
+        initialize_model(
+          model_name: model_name,
+          cache_dir: cache_dir,
+          threads: threads,
+          providers: providers,
+          show_progress: show_progress,
+          quantization: quantization
+        )
+      end
+
+      setup_model_and_tokenizer(model_file_override: model_file || quantized_model_file)
+    end
+
+    # Generate sparse embeddings for documents
+    #
+    # @param documents [Array<String>, String] Text document(s) to embed
+    # @param batch_size [Integer] Number of documents to process at once
+    # @yield [Progress] Optional progress callback called after each batch
+    # @return [Enumerator] Lazy enumerator yielding SparseEmbedding objects
+    def embed(documents, batch_size: 32, &progress_callback)
+      documents = Validators.validate_documents!(documents)
+      return Enumerator.new { |_| } if documents.empty?
+
+      total_batches = (documents.length.to_f / batch_size).ceil
+
+      Enumerator.new do |yielder|
+        documents.each_slice(batch_size).with_index(1) do |batch, batch_num|
+          sparse_embeddings = compute_sparse_embeddings(batch)
+          sparse_embeddings.each { |emb| yielder << emb }
+
+          if progress_callback
+            progress = Progress.new(current: batch_num, total: total_batches, batch_size: batch_size)
+            progress_callback.call(progress)
+          end
+        end
+      end
+    end
+
+    # Generate sparse embeddings for queries (with "query: " prefix)
+    #
+    # @param queries [Array<String>, String] Query text(s) to embed
+    # @param batch_size [Integer] Number of queries to process at once
+    # @return [Enumerator] Lazy enumerator yielding SparseEmbedding objects
+    def query_embed(queries, batch_size: 32)
+      queries = [queries] if queries.is_a?(String)
+      prefixed = queries.map { |q| "query: #{q}" }
+      embed(prefixed, batch_size: batch_size)
+    end
+
+    # Generate sparse embeddings for passages
+    #
+    # @param passages [Array<String>, String] Passage text(s) to embed
+    # @param batch_size [Integer] Number of passages to process at once
+    # @return [Enumerator] Lazy enumerator yielding SparseEmbedding objects
+    def passage_embed(passages, batch_size: 32)
+      passages = [passages] if passages.is_a?(String)
+      embed(passages, batch_size: batch_size)
+    end
+
+    # Generate sparse embeddings asynchronously
+    #
+    # @param documents [Array<String>, String] Text document(s) to embed
+    # @param batch_size [Integer] Number of documents to process at once
+    # @return [Async::Future] Future that resolves to array of SparseEmbedding objects
+    def embed_async(documents, batch_size: 32)
+      Async::Future.new { embed(documents, batch_size: batch_size).to_a }
+    end
+
+    # Generate query embeddings asynchronously
+    #
+    # @param queries [Array<String>, String] Query text(s) to embed
+    # @param batch_size [Integer] Number of queries to process at once
+    # @return [Async::Future] Future that resolves to array of SparseEmbedding objects
+    def query_embed_async(queries, batch_size: 32)
+      Async::Future.new { query_embed(queries, batch_size: batch_size).to_a }
+    end
+
+    # Generate passage embeddings asynchronously
+    #
+    # @param passages [Array<String>, String] Passage text(s) to embed
+    # @param batch_size [Integer] Number of passages to process at once
+    # @return [Async::Future] Future that resolves to array of SparseEmbedding objects
+    def passage_embed_async(passages, batch_size: 32)
+      Async::Future.new { passage_embed(passages, batch_size: batch_size).to_a }
+    end
+
+    # List all supported sparse models
+    #
+    # @return [Array<Hash>] Array of model information hashes
+    def self.list_supported_models
+      SUPPORTED_SPARSE_MODELS.values.map(&:to_h)
+    end
+
+    private
+
+    def resolve_model_info(model_name)
+      # Check built-in registry first
+      info = SUPPORTED_SPARSE_MODELS[model_name]
+      return info if info
+
+      # Check custom registry
+      info = CustomModelRegistry.sparse_models[model_name]
+      return info if info
+
+      raise Error, "Unknown sparse model: #{model_name}"
+    end
+
+    def create_local_model_info(model_name:, model_file:, tokenizer_file:)
+      SparseModelInfo.new(
+        model_name: model_name,
+        description: 'Local sparse model',
+        size_in_gb: 0,
+        sources: {},
+        model_file: model_file || 'model.onnx',
+        tokenizer_file: tokenizer_file || 'tokenizer.json'
+      )
+    end
+
+    def compute_sparse_embeddings(texts)
+      prepared = tokenize_and_prepare(texts)
+      outputs = @session.run(nil, prepared[:inputs])
+      logits = extract_logits(outputs)
+
+      # Convert to sparse embeddings
+      texts.length.times.map do |i|
+        create_sparse_embedding(logits[i], prepared[:attention_mask][i])
+      end
+    end
+
+    def extract_logits(outputs)
+      # SPLADE outputs logits for each token position
+      if outputs.is_a?(Hash)
+        outputs['logits'] || outputs.values.first
+      else
+        outputs.first
+      end
+    end
+
+    # Transform token logits into sparse embedding using SPLADE algorithm
+    #
+    # SPLADE (Sparse Lexical AnD Expansion) uses a log-saturation function:
+    # weight = log(1 + ReLU(logit))
+    #
+    # This ensures:
+    # - ReLU removes negative activations (irrelevant terms)
+    # - log1p provides saturation to prevent any term from dominating
+    # - Max-pooling across positions captures the strongest signal per vocabulary term
+    #
+    # @param token_logits [Array<Array<Float>>] Logits for each token position
+    # @param attention_mask [Array<Integer>] Mask indicating valid positions
+    # @return [SparseEmbedding] Sparse vector with vocabulary indices and weights
+    def create_sparse_embedding(token_logits, attention_mask)
+      vocab_size = token_logits.first.length
+      max_weights = Array.new(vocab_size, 0.0)
+
+      token_logits.each_with_index do |logits, pos|
+        next if attention_mask[pos].zero?
+
+        logits.each_with_index do |logit, vocab_idx|
+          # SPLADE transformation: log(1 + ReLU(x))
+          activated = logit.positive? ? logit : 0.0
+          weight = Math.log(1.0 + activated)
+          max_weights[vocab_idx] = weight if weight > max_weights[vocab_idx]
+        end
+      end
+
+      # Extract non-zero indices and values
+      indices = []
+      values = []
+
+      max_weights.each_with_index do |weight, idx|
+        if weight.positive?
+          indices << idx
+          values << weight
+        end
+      end
+
+      SparseEmbedding.new(indices: indices, values: values)
+    end
+  end
+end

data/lib/fastembed/sparse_model_info.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module Fastembed
+  # Model information for sparse embedding models (SPLADE, etc.)
+  #
+  # Sparse models like SPLADE produce vectors where most dimensions are zero,
+  # with non-zero values corresponding to vocabulary tokens. These are useful
+  # for hybrid search combining with dense embeddings.
+  #
+  # @example Access sparse model info
+  #   info = Fastembed::SUPPORTED_SPARSE_MODELS['prithivida/Splade_PP_en_v1']
+  #   info.model_name  # => "prithivida/Splade_PP_en_v1"
+  #
+  class SparseModelInfo
+    include BaseModelInfo
+
+    # Create a new SparseModelInfo instance
+    #
+    # @param model_name [String] Full model identifier
+    # @param description [String] Human-readable description
+    # @param size_in_gb [Float] Model size in GB
+    # @param sources [Hash] Source repositories
+    # @param model_file [String] Path to ONNX model file
+    # @param tokenizer_file [String] Path to tokenizer file
+    # @param max_length [Integer] Maximum sequence length
+    def initialize(
+      model_name:,
+      description:,
+      size_in_gb:,
+      sources:,
+      model_file: 'onnx/model.onnx',
+      tokenizer_file: 'tokenizer.json',
+      max_length: BaseModelInfo::DEFAULT_MAX_LENGTH
+    )
+      initialize_base(
+        model_name: model_name,
+        description: description,
+        size_in_gb: size_in_gb,
+        sources: sources,
+        model_file: model_file,
+        tokenizer_file: tokenizer_file,
+        max_length: max_length
+      )
+    end
+
+    # Convert to hash representation
+    # @return [Hash] Model info as a hash
+    def to_h
+      {
+        model_name: model_name,
+        description: description,
+        size_in_gb: size_in_gb,
+        sources: sources,
+        model_file: model_file,
+        tokenizer_file: tokenizer_file,
+        max_length: max_length
+      }
+    end
+  end
+
+  # Registry of supported sparse embedding models
+  SUPPORTED_SPARSE_MODELS = {
+    'prithivida/Splade_PP_en_v1' => SparseModelInfo.new(
+      model_name: 'prithivida/Splade_PP_en_v1',
+      description: 'SPLADE++ model for sparse text retrieval',
+      size_in_gb: 0.53,
+      sources: { hf: 'prithivida/Splade_PP_en_v1' }
+      # Uses default model_file: 'onnx/model.onnx'
+    ),
+    'prithivida/Splade_PP_en_v2' => SparseModelInfo.new(
+      model_name: 'prithivida/Splade_PP_en_v2',
+      description: 'SPLADE++ v2 with improved performance',
+      size_in_gb: 0.53,
+      sources: { hf: 'prithivida/Splade_PP_en_v2' }
+      # Uses default model_file: 'onnx/model.onnx'
+    )
+  }.freeze
+
+  DEFAULT_SPARSE_MODEL = 'prithivida/Splade_PP_en_v1'
+end