Generative AI Google Brain

search-tuning-vertexai-search-tuning

4 min read

Search tuning in Vertex AI Search

Source notebook

Repo path: search/tuning/vertexai-search-tuning.ipynb · Open on GitHub · advanced

Tunes Vertex AI Search with JSONL/TSV Q&A data and tests a search app over Cloud Storage PDFs.

Summary

This notebook demonstrates how to prepare search tuning data from Kubernetes FAQ Markdown files, generating JSONL corpus/query files and TSV train/test mappings. It creates a Cloud Storage-backed Vertex AI Search data store, imports PDF documents, builds a search app with LLM summaries and citations, then submits a Discovery Engine search-tuning job. The workflow also shows pre-tuning search testing and stresses that the search app must be rebuilt after datastore refreshes.

Key code patterns

Initialize Vertex AI context

PROJECT_ID = "genai-customersupport"
LOCATION = "global"
STORAGE_LOCATION = "us"
 
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the project and global location used by Vertex AI and Discovery Engine resources.

Create JSONL corpus and query files

json_line = '{{"_id": "ans{:04d}", "text": "{}" }}\n'.format(
    idx, str.strip(answer).replace("\n", " ")
)
 
json_line = '{{ "_id": "que{:04d}", "text": "{}" }}\n'.format(
    idx, str.strip(question).replace("\n", " ")
)

Search tuning data uses identifier-text JSONL records for answer corpus and user queries.

Create TSV train and test mappings

json_line = "query-id\tcorpus-id\tscore\n"
for question in questions:
    json_line = f"que{idx:04d}\tans{idx:04d}\t1\n"
    jsonfile += json_line
    if idx > 0.85 * len_questions:
        break

Maps query IDs to corpus IDs with relevance scores and splits roughly 85 percent for training.

Create Vertex AI Search data store

data_store = discoveryengine.DataStore(
    display_name=data_store_name,
    industry_vertical=discoveryengine.IndustryVertical.GENERIC,
    solution_types=[discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH],
    content_config=discoveryengine.DataStore.ContentConfig.CONTENT_REQUIRED,
)
operation = client.create_data_store(request=request)

Defines a generic search data store that requires document content for indexing.

Import Cloud Storage documents

document_import_request = discoveryengine.ImportDocumentsRequest(
    parent=parent,
    gcs_source=discoveryengine.GcsSource(
        input_uris=[f"{SEARCH_DATASTORE_PATH_REMOTE}/*"],
        data_schema="content",
    ),
    reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
)

Indexes unstructured documents from Cloud Storage using incremental reconciliation.

Create LLM-enabled search app

engine = discoveryengine.Engine(
    display_name=engine_name,
    industry_vertical=discoveryengine.IndustryVertical.GENERIC,
    solution_type=discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH,
    search_engine_config=discoveryengine.Engine.SearchEngineConfig(
        search_tier=discoveryengine.SearchTier.SEARCH_TIER_ENTERPRISE,
        search_add_ons=[discoveryengine.SearchAddOn.SEARCH_ADD_ON_LLM],
    ),
    data_store_ids=data_store_ids,
)

Builds an enterprise Vertex AI Search app with the LLM add-on for generated summaries.

Search with snippets and citations

content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec(
    snippet_spec=discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(
        return_snippet=True
    ),
    summary_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec(
        summary_result_count=5,
        include_citations=True,
        model_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec.ModelSpec(version="stable"),
    ),
)

Requests snippets and cited summaries from search results.

Submit search tuning job

operation = client.train_custom_model(
    request=discoveryengine.TrainCustomModelRequest(
        gcs_training_input=discoveryengine.TrainCustomModelRequest.GcsTrainingInput(
            corpus_data_path=corpus_data_path,
            query_data_path=query_data_path,
            train_data_path=train_data_path,
            test_data_path=test_data_path,
        ),
        data_store=data_store,
        model_type="search-tuning",
    )
)

Starts a Discovery Engine custom model training operation for Vertex AI Search tuning.

Models & APIs used

  • APIs / services: Vertex AI, Vertex AI Search, Discovery Engine, Cloud Storage
  • SDKs / libraries: vertexai, google-cloud-discoveryengine, google-cloud-aiplatform, shortuuid

When to use this

Use this pattern when you need to tune Vertex AI Search responses with domain-specific Q&A data and indexed documents.

Gotchas & caveats

  • The Vertex AI API must be enabled and the notebook authenticates with gcloud auth login or Colab auth.
  • The runtime must restart after installing google-cloud-aiplatform and google-cloud-discoveryengine.
  • The notebook sets LOCATION to global and only builds regional Discovery Engine client options when LOCATION is not global.
  • Tuning data must be uploaded to Cloud Storage as JSONL corpus/query files and TSV train/test files.
  • Markdown FAQ files are converted to PDF because the notebook states Vertex AI Search cannot accept Markdown files.
  • PDF conversion is platform-specific: macOS uses Homebrew and xelatex, Linux uses apt-get and pdflatex.
  • BEIR and SciFact datasets are described as large enough to make tuning jobs run too long, so the notebook starts with small Kubernetes FAQ data.
  • The notebook states users must rebuild the search app after rebuilding or refreshing the datastore.

Best practices

  • Prepare datasets in JSONL format with identifier-text pairs.
  • Represent query-to-answer mappings in tab-separated TSV files.
  • Upload additional tuning documents and datasets to Cloud Storage before datastore refresh and tuning.
  • Use INCREMENTAL document import when adding files to the datastore.
  • Keep a baseline search result before tuning so it can be compared with post-tuning behavior.
  • Configure summaries with citations when checking correctness against source documents.
  • Use generated PDFs for Markdown source documents before importing into the Vertex AI Search datastore.
  • Delete engines, purge documents, and delete data stores with helper functions when cleaning up resources.

Graph View

Backlinks