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 



761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
 
# File 'lib/informers/pipelines.rb', line 761

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