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 



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
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
 
# File 'lib/informers/pipelines.rb', line 822

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 previous revision of 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