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 



247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
 
# File 'lib/informers/pipelines.rb', line 247

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