Class: Informers::FeatureExtractionPipeline

Inherits:
Pipeline
  • Object
show all
Defined in:
lib/informers/pipelines.rb

Direct Known Subclasses

EmbeddingPipeline

Instance Method Summary collapse

Methods inherited from Pipeline

#initialize

Constructor Details

This class inherits a constructor from Informers::Pipeline

Instance Method Details

#call(texts, pooling: "none", normalize: false, quantize: false, precision: "binary", model_output: nil) ⇒ Object 



788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
 
# File 'lib/informers/pipelines.rb', line 788

def call(
  texts,
  pooling: "none",
  normalize: false,
  quantize: false,
  precision: "binary",
  model_output: nil
)
  # Run tokenization
  model_inputs = @tokenizer.(texts,
    padding: true,
    truncation: true
  )
  model_options = {}

  if !model_output.nil?
    model_options[:output_names] = Array(model_output)
  elsif @model.instance_variable_get(:@output_names) == ["token_embeddings"] && pooling == "mean" && normalize
    # optimization for sentence-transformers/all-MiniLM-L6-v2
    model_options[:output_names] = ["sentence_embedding"]
    pooling = "none"
    normalize = false
  end

  # Run model
  outputs = @model.(model_inputs, **model_options)

  # TODO improve
  result =
    if outputs.is_a?(Array)
      # TODO show returned instead of all
      output_names = @model.instance_variable_get(:@session).outputs.map { |v| v[:name] }
      raise Error, "unexpected outputs: #{output_names}" if outputs.size != 1
      outputs[0]
    else
      outputs.logits
    end

  case pooling
  when "none"
    # Skip pooling
  when "mean"
    result = Utils.mean_pooling(result, model_inputs[:attention_mask])
  when "cls"
    result = result.map(&:first)
  else
    # TODO raise ArgumentError in 2.0
    raise Error, "Pooling method '#{pooling}' not supported."
  end

  if normalize
    result = Utils.normalize(result)
  end

  if quantize
    result = quantize_embeddings(result, precision)
  end

  texts.is_a?(Array) ? result : result[0]
end