Fix PaliGemma conversion and verify.py of examples broken by ExportCo…

…nfig.

PiperOrigin-RevId: 704971311

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -15,6 +15,8 @@
 
 """Example of building a decoder of PaliGemma 3B model which is Gemma1."""
 
+from typing import Optional
+
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
@@ -51,6 +53,7 @@ def forward(
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       input_embeds: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(tokens, input_pos, kv_cache)

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -16,11 +16,13 @@
 """Example of building a full-stack of PaliGemma model."""
 
 from dataclasses import dataclass
+from typing import Optional
 
 from ai_edge_torch.generative.examples.paligemma import decoder
 from ai_edge_torch.generative.examples.paligemma import image_encoder
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -67,9 +69,16 @@ def forward(
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       pixel_values: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if pixel_values is None:
-      return self.decoder(tokens, input_pos, kv_cache)
+      return self.decoder(
+          tokens=tokens,
+          input_pos=input_pos,
+          kv_cache=kv_cache,
+          input_embeds=None,
+          export_config=export_config
+      )
 
     input_embeds = self.decoder.tok_embedding(tokens)
 
@@ -100,6 +109,7 @@ def forward(
         input_pos=input_pos,
         kv_cache=kv_cache,
         input_embeds=input_embeds,
+        export_config=export_config,
     )
 
 

ai_edge_torch/generative/utilities/transformers_verifier.py

@@ -29,7 +29,7 @@ class TransformersModelWrapper(verifier.ModelWrapper):
   an object with `logits` field.
 
   Transformers models get `max_new_tokens` settings for generate() via
-  ExportConfig.
+  GenerationConfig.
   """
 
   def forward(self, tokens: torch.Tensor) -> torch.Tensor:
@@ -38,5 +38,5 @@ def forward(self, tokens: torch.Tensor) -> torch.Tensor:
   def generate(
       self, inputs: torch.Tensor, max_new_tokens: int
   ) -> torch.IntTensor:
-    export_config = transformers.ExportConfig(max_new_tokens=max_new_tokens)
-    return self.model.generate(inputs=inputs, generation_config=export_config)
+    gen_config = transformers.GenerationConfig(max_new_tokens=max_new_tokens)
+    return self.model.generate(inputs=inputs, generation_config=gen_config)

ai_edge_torch/generative/utilities/verifier.py

@@ -115,7 +115,7 @@ def _forward_with_kv_cache(
     # pixel_values only when it is not None. Otherwise, it may raise an error.
     if pixel_values is None:
       output = self.model.forward(
-          tokens, input_pos, kv_cache, self.export_config
+          tokens, input_pos, kv_cache, export_config=self.export_config
       )
     else:
       output = self.model.forward(