Generative AI Google Brain

gemini-use-cases-applying-llms-to-data-gemini-and-documentai-for-entity-extraction-test-document-ai-gemini

2 min read

Test Document AI Gemini

Source notebook

Repo path: gemini/use-cases/applying-llms-to-data/gemini-and-documentai-for-entity-extraction/test_document_ai_gemini.ipynb · Open on GitHub · intermediate

Compares Document AI entity extraction with Gemini-based extraction on a PDF.

Summary

This notebook processes a local PDF with a Document AI processor and extracts entities from the returned document. It then uploads the same PDF to Cloud Storage, uses Gemini 2.0 Flash with an extraction prompt to extract entities, and asks Gemini to compare the Document AI and Gemini outputs.

Diagrams

diagram.png — source

Key code patterns

Document AI online extraction

online_extractor = OnlineDocumentExtractor(
    project_id=project_id,
    location=location,
    processor_id=processor_id,
)
online_document = online_extractor.process_document(file_path, mime_type)
docai_entities = DocumentAIEntityExtractor(online_document).extract_entities()

Runs a configured Document AI processor on a PDF and extracts processor entities.

Gemini extraction from GCS file

temp_file_uploader = TempFileUploader(gcs_temp_uri)
gcs_input_uri = temp_file_uploader.upload_file(file_path)
model_extractor = ModelBasedEntityExtractor(
    "gemini-2.0-flash", prompt_extract, gcs_input_uri
)
gemini_entities = model_extractor.extract_entities()
temp_file_uploader.delete_file()

Uploads the PDF to temporary Cloud Storage and uses Gemini with an extraction prompt.

Gemini comparison analysis

compare_prompt = get_compare_entities_prompt().format(
    docai_output=str(docai_entities),
    gemini_output=str(gemini_entities),
)
model = GenerativeModel("gemini-2.0-flash")
response = model.generate_content(compare_prompt)
print(response.text)

Uses Gemini to compare the two entity extraction results in one generated analysis.

Models & APIs used

  • Models: gemini-2.0-flash
  • APIs / services: Document AI, Vertex AI, Cloud Storage
  • SDKs / libraries: vertexai

When to use this

Use this pattern to evaluate Document AI entity extraction against Gemini extraction for the same PDF document.

Gotchas & caveats

  • Requires valid project_id, location, processor_id, and a supported processor location such as us or eu.
  • The PDF must be available locally as file_path with the correct mime_type.
  • Gemini extraction path requires a writable temporary Cloud Storage URI.
  • Temporary uploaded files should be deleted after extraction.
  • processor_version_id is shown as optional for batch processing but not used in the online example.

Best practices

  • Process the same input document through both extraction paths before comparing outputs.
  • Use separate prompts for extraction and comparison.
  • Upload the document to temporary Cloud Storage for model-based extraction, then delete the temporary file.
  • Keep Document AI configuration values explicit: project, location, processor, file path, and MIME type.

Graph View

Backlinks