Generative AI Google Brain

open-models-evaluation-vertex-ai-tgi-gemma-with-genai-evaluation

4 min read

Hugging Face DLCs: Using Gemma for running evaluations with Vertex AI Gen AI Evaluation

Source notebook

Repo path: open-models/evaluation/vertex_ai_tgi_gemma_with_genai_evaluation.ipynb · Open on GitHub · advanced

Deploys Gemma 2 on Vertex AI TGI and evaluates summarization with Gen AI Evaluation.

Summary

The notebook shows how to deploy the gated Hugging Face model google/gemma-2-9b-it to a Vertex AI endpoint using a Hugging Face Text Generation Inference container. It evaluates Gemma 2 on an XSum summarization task with EvalTask, using rouge_l_sum, summarization_quality, and fluency metrics. It then generates summaries with gemini-2.0-flash and uses Gemma 2 as a custom LLM-as-judge metric for catchiness.

Key code patterns

Deploy Gemma on TGI

gemma_model = aiplatform.Model.upload(
    display_name="google--gemma-2-9b-it",
    serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-3.ubuntu2204.py311",
    serving_container_environment_variables={"MODEL_ID": "google/gemma-2-9b-it", "HUGGING_FACE_HUB_TOKEN": get_token()},
    serving_container_ports=[8080],
)
deployed_gemma_model = gemma_model.deploy(
    endpoint=aiplatform.Endpoint.create(display_name="google--gemma-2-9b-it-endpoint"),
    machine_type="g2-standard-24", accelerator_type="NVIDIA_L4", accelerator_count=2,
)

Registers a Hugging Face DLC model and deploys it to a GPU-backed Vertex AI endpoint.

Wrap endpoint for EvalTask

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it")
 
def gemma_fn(prompt, generation_config=generation_config):
    formatted_prompt = tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True
    )
    instance = {"inputs": formatted_prompt, "parameters": generation_config}
    output = deployed_gemma_model.predict(instances=[instance])
    return output.predictions[0]

Turns the deployed TGI endpoint into a callable model function compatible with EvalTask.

Evaluate summarization

prompt_template = "Summarize the following article in one sentence: {context}.\nSummary:"
metrics = ["rouge_l_sum", "summarization_quality", "fluency"]
 
eval_task = EvalTask(
    dataset=eval_model_sample_df,
    metrics=metrics,
    experiment="eval-gemma-base-prompt-sum",
)
eval_result = eval_task.evaluate(
    model=gemma_fn,
    prompt_template=prompt_template,
    experiment_run_name=experiment_run_name,
)

Runs prompt-template evaluation with built-in summarization and fluency metrics.

Generate Gemini responses

eval_model_sample_df["prompt"] = eval_model_sample_df.apply(
    lambda row: prompt_template.format(context=row["context"]), axis=1
)
gemini_llm = init_new_model(model_name="gemini-2.0-flash")
gemini_predictions = [
    async_generate(p, model=gemini_llm) for p in eval_model_sample_df["prompt"]
]
gemini_predictions_col = await tqdm_asyncio.gather(*gemini_predictions)
eval_model_sample_df["response"] = gemini_predictions_col

Creates generated summaries to be scored later by Gemma as an evaluator.

Custom LLM judge metric

def catchiness_fn(instance: dict) -> dict:
    metric_prompt = metric_prompt_template.format(
        prompt=instance["prompt"], response=instance["response"]
    )
    eval_response = gemma_fn(metric_prompt, {"max_new_tokens": 256, "temperature": 0})
    return parse_json_output(eval_response)
 
catchiness_metric = CustomMetric(
    name="catchiness",
    metric_function=catchiness_fn,
)

Uses Gemma 2 as an autorater by returning a score and explanation from a custom metric function.

Models & APIs used

  • Models: google/gemma-2-9b-it, gemini-2.0-flash
  • APIs / services: Vertex AI, Vertex AI Model Garden, Vertex AI Prediction, Vertex AI Model Eval, Vertex AI Gen AI Evaluation service, Vertex AI Experiments, Gemini API on Vertex AI, Artifact Registry
  • SDKs / libraries: google-cloud-aiplatform, vertexai, datasets, huggingface_hub, transformers, pandas, plotly, tenacity

When to use this

Use this pattern when evaluating a Hugging Face Gemma model deployed on Vertex AI, including model-as-target and model-as-judge workflows.

Gotchas & caveats

  • google/gemma-2-9b-it is gated, so the user must accept the Hugging Face license and provide a read-only Hugging Face token.
  • The notebook requires an existing Google Cloud project with aiplatform.googleapis.com and artifactregistry.googleapis.com enabled.
  • The compute service account is granted roles/aiplatform.user before Vertex AI deployment and evaluation.
  • Colab users must restart the runtime after package installation and authenticate with google.colab.auth.
  • The deployment requests g2-standard-24 with two NVIDIA_L4 accelerators, so region availability and quota must support that shape.
  • The sample size defaults to n=10, so reported metrics are lightweight demonstration results unless expanded.
  • The custom catchiness parser returns catchiness 0 and an empty explanation when the judge output is not valid JSON.

Best practices

  • Initialize Vertex AI with project and location before creating models, endpoints, and evaluations.
  • Use a tokenizer chat template before sending prompts to the Hugging Face TGI endpoint.
  • Wrap endpoint prediction in a model function so EvalTask can call it consistently.
  • Track evaluation runs with Vertex AI Experiments through experiment and experiment_run_name.
  • Combine reference-based metrics such as rouge_l_sum with model-based metrics such as summarization_quality and fluency.
  • Use retry with wait_random_exponential for asynchronous Gemini generation calls.
  • Filter and sample the evaluation dataset before running a small tutorial evaluation.
  • Provide optional cleanup logic for created Vertex AI experiments.

Graph View

Backlinks