Class: Informers::DocumentQuestionAnsweringPipeline

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

Instance Method Summary collapse

Methods inherited from Pipeline

#initialize

Constructor Details

This class inherits a constructor from Informers::Pipeline

Instance Method Details

#call(image, question, **generate_kwargs) ⇒ Object 



747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
 
# File 'lib/informers/pipelines.rb', line 747

def call(image, question, **generate_kwargs)
  # NOTE: For now, we only support a batch size of 1

  # Preprocess image
  prepared_image = prepare_images(image)[0]
  pixel_values = @processor.(prepared_image)[:pixel_values]

  # Run tokenization
  task_prompt = "<s_docvqa><s_question>#{question}</s_question><s_answer>"
  decoder_input_ids =
    @tokenizer.(
      task_prompt,
      add_special_tokens: false,
      padding: true,
      truncation: true
    )[:input_ids]

  # Run model
  output =
    @model.generate(
      pixel_values,
      generate_kwargs.merge(
        decoder_input_ids: decoder_input_ids[0],
        max_length: @model.config["decoder"]["max_position_embeddings"]
      ).transform_keys(&:to_s)
    )

  # Decode output
  decoded = @tokenizer.batch_decode(output, skip_special_tokens: false)[0]

  # Parse answer
  match = decoded.match(/<s_answer>(.*?)<\/s_answer>/)
  answer = nil
  if match && match.length >= 2
    answer = match[1].strip
  end
  [{answer:}]
end