ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -0,0 +1,160 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Example of building an image encoder of PaliGemma model which is Siglip."""
+
+from ai_edge_torch.generative.layers import attention
+from ai_edge_torch.generative.layers import builder
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="vision_tower.vision_model.encoder.layers.{}.mlp.fc1",
+    ff_down_proj="vision_tower.vision_model.encoder.layers.{}.mlp.fc2",
+    attn_query_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.q_proj"
+    ),
+    attn_key_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.k_proj"
+    ),
+    attn_value_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.v_proj"
+    ),
+    attn_output_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.out_proj"
+    ),
+    pre_attn_norm="vision_tower.vision_model.encoder.layers.{}.layer_norm1",
+    post_attn_norm="vision_tower.vision_model.encoder.layers.{}.layer_norm2",
+    embedding="vision_tower.vision_model.embeddings.patch_embedding",
+    embedding_position=(
+        "vision_tower.vision_model.embeddings.position_embedding.weight"
+    ),
+    final_norm="vision_tower.vision_model.post_layernorm",
+)
+
+
+class SiglipVisionEncoder(nn.Module):
+  """Signlip vision encoder from the Edge Generative API."""
+
+  def __init__(self, config: cfg.ModelConfig):
+    super().__init__()
+
+    # Construct model layers.
+    self.tok_embedding = nn.Conv2d(
+        in_channels=config.image_embedding.channels,
+        out_channels=config.embedding_dim,
+        kernel_size=config.image_embedding.patch_size,
+        stride=config.image_embedding.patch_size,
+        padding=0,
+    )
+    num_patches = (
+        config.image_embedding.image_size // config.image_embedding.patch_size
+    ) ** 2
+    self.tok_embedding_position = nn.Parameter(
+        torch.zeros((num_patches, config.embedding_dim))
+    )
+
+    self.transformer_blocks = nn.ModuleList(
+        attention.TransformerBlock(config.block_config(idx), config)
+        for idx in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+    # Embed the image according to SiplipVisionEmbeddings.
+    x = self.tok_embedding(pixel_values)
+    x = x.flatten(2).transpose(1, 2) + self.tok_embedding_position
+
+    # Pass a dummy mask because SDPA attention impl expects non-None mask.
+    mask = torch.zeros(x.shape[:2])
+    for _, block in enumerate(self.transformer_blocks):
+      x = block(x, mask=mask)
+    return self.final_norm(x)
+
+
+def get_image_encoder_config() -> cfg.ModelConfig:
+  """Returns the model config for the image encoder of a PaliGemma 3B-224 model.
+
+  Returns:
+    The model config for the image encoder of a PaliGemma 3B model.
+  """
+  image_embedding_config = cfg.ImageEmbeddingConfig(
+      channels=3,
+      image_size=224,
+      patch_size=14,
+  )
+  attn_config = cfg.AttentionConfig(
+      num_heads=16,
+      head_dim=72,
+      num_query_groups=16,
+      qkv_use_bias=True,
+      output_proj_use_bias=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=4304,
+      use_bias=True,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.LAYER_NORM,
+      epsilon=1e-6,
+      enable_hlfb=True,
+  )
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=0,  # Not used in image encoder.
+      num_layers=27,
+      max_seq_len=0,  # Not used in image encoder.
+      embedding_dim=1152,
+      embedding_use_bias=True,
+      image_embedding=image_embedding_config,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_fake_image_encoder_config() -> cfg.ModelConfig:
+  config = get_image_encoder_config()
+  # PaliGemma image encoder has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 128
+  config.image_embedding.image_size = 8
+  config.image_embedding.patch_size = 2
+  config.num_layers = 2
+  return config
+
+
+def build_image_encoder(checkpoint_path: str) -> SiglipVisionEncoder:
+  config = get_image_encoder_config()
+  encoder = SiglipVisionEncoder(config)
+  loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
+  # Loose the strictness because only image encoder is being loaded.
+  loader.load(encoder, strict=False)
+  encoder.eval()
+  return encoder

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -0,0 +1,179 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""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 decoder2
+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
+
+PROJECTION_TENSOR_NAME = "multi_modal_projector.linear"
+
+
+@dataclass
+class PaliGemmaConfig:
+  """PaliGemma model configurations."""
+
+  image_encoder_config: cfg.ModelConfig
+  decoder_config: cfg.ModelConfig
+
+  image_token_id: int
+  image_projection_scale: float
+  image_projection_use_bias: bool = False
+
+
+class PaliGemma(nn.Module):
+  """PaliGemma model from the Edge Generative API."""
+
+  def __init__(self, config: PaliGemmaConfig, decoder_class: nn.Module):
+    super().__init__()
+
+    self.image_encoder = image_encoder.SiglipVisionEncoder(
+        config.image_encoder_config
+    )
+    self.image_projection = nn.Linear(
+        config.image_encoder_config.embedding_dim,
+        config.decoder_config.embedding_dim,
+        bias=config.image_projection_use_bias,
+    )
+    self.decoder = decoder_class(config.decoder_config)
+    image_embedding_config = config.image_encoder_config.image_embedding
+    self.num_patches = (
+        image_embedding_config.image_size // image_embedding_config.patch_size
+    ) ** 2
+    self.config = config
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
+      pixel_values: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+      called_by_generate: bool = True,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    if pixel_values is None:
+      return self.decoder(
+          tokens=tokens,
+          input_pos=input_pos,
+          kv_cache=kv_cache,
+          mask=mask,
+          input_embeds=None,
+          export_config=export_config,
+          called_by_generate=called_by_generate,
+      )
+
+    input_embeds = self.decoder.tok_embedding(tokens)
+
+    image_encoded = self.image_encoder(pixel_values=pixel_values)
+    image_embeds = self.image_projection(image_encoded)
+    image_embeds = image_embeds / self.config.image_projection_scale
+
+    # Merging image_embeds into text_embeds as PaliGemmaForConditionalGeneration
+    # can be done like:
+    #
+    #   image_mask = tokens == self.config.image_token_id
+    #   image_mask = image_mask.unsqueeze(-1).expand_as(input_embeds)
+    #   input_embeds = input_embeds.masked_scatter(image_mask, image_embeds)
+    #
+    # Unfortunately, torch.Tensor.masked_scatter can't be lowered on CPU.
+    # Since PaliGemma token embedder reserves the first [num_patches] tokens
+    # for image tokens, we can use this property to merge image_embeds into
+    # input_embeds by concatenating them.
+    assert image_embeds.shape[1] == self.num_patches
+    assert input_embeds.shape[1] >= self.num_patches
+    input_embeds = torch.cat(
+        (image_embeds, input_embeds[:, self.num_patches:, :]), dim=1
+    )
+
+    return self.decoder(
+        tokens=None,
+        input_pos=input_pos,
+        kv_cache=kv_cache,
+        mask=mask,
+        input_embeds=input_embeds,
+        export_config=export_config,
+        called_by_generate=called_by_generate,
+    )
+
+
+def get_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
+  """Returns the model config for a PaliGemma 3B-224 model.
+
+  Returns:
+    The model config for a PaliGemma 3B model.
+  """
+  return PaliGemmaConfig(
+      image_encoder_config=image_encoder.get_image_encoder_config(),
+      decoder_config=get_decoder_config(**kwargs),
+      image_token_id=257152,
+      image_projection_scale=2048**0.5,
+      image_projection_use_bias=True,
+  )
+
+
+def get_fake_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
+  return PaliGemmaConfig(
+      image_encoder_config=image_encoder.get_fake_image_encoder_config(),
+      decoder_config=get_decoder_config(**kwargs),
+      image_token_id=127,
+      image_projection_scale=128**0.5,
+      image_projection_use_bias=True,
+  )
+
+
+def build_model(checkpoint_path: str, version: int = 2, **kwargs) -> PaliGemma:
+  if version == 1:
+    decoder_class = decoder.Decoder
+    decoder_tensor_names = decoder.TENSOR_NAMES
+    get_decoder_config = decoder.get_decoder_config
+  else:
+    decoder_class = decoder2.Decoder2
+    decoder_tensor_names = decoder2.TENSOR_NAMES
+    get_decoder_config = decoder2.get_decoder2_config
+
+  config = get_model_config(get_decoder_config, **kwargs)
+  model = PaliGemma(config, decoder_class)
+  # Load the parameters of image encoder.
+  loader = loading_utils.ModelLoader(
+      checkpoint_path, image_encoder.TENSOR_NAMES
+  )
+  loader.load(model.image_encoder, strict=False)
+  # Load the parameters of decoder.
+  loader = loading_utils.ModelLoader(checkpoint_path, decoder_tensor_names)
+  loader.load(model.decoder, strict=False)
+
+  # Load the parameters of image projection.
+  loader = loading_utils.ModelLoader(checkpoint_path, None)
+  state = loader.get_state()
+  converted_state = dict()
+  converted_state["weight"] = state.pop(f"{PROJECTION_TENSOR_NAME}.weight")
+  if config.image_projection_use_bias:
+    converted_state["bias"] = state.pop(f"{PROJECTION_TENSOR_NAME}.bias")
+  model.image_projection.load_state_dict(converted_state)
+
+  model.eval()
+  return model

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -0,0 +1,161 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Verifies the reauthored PaliGemma 3B model."""
+
+import logging
+import pathlib
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.paligemma import paligemma
+from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.utilities import verifier
+import kagglehub
+from PIL import Image
+import requests
+import torch
+import transformers
+
+_VERSION = flags.DEFINE_enum(
+    "version",
+    "2",
+    ["1", "2"],
+    "The version of PaliGemma model to verify.",
+)
+_IMAGE_URL = flags.DEFINE_string(
+    "image_url",
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
+    "The image URI to encode.",
+)
+_PROMPTS = flags.DEFINE_string(
+    "prompts",
+    "describe en",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+_CHECKPOINT = {
+    "1": "google/paligemma-3b-mix-224",
+    "2": "google/paligemma-2/transformers/paligemma2-3b-pt-224",
+}
+
+
+class ReauthoredPaliGemmaWrapper(verifier.ReauthoredModelWrapper):
+  """Reauthored PaliGemma model wrapper."""
+
+  def __init__(self, model: torch.nn.Module):
+    super().__init__(model)
+    self.forward_called_by_generate = False
+
+  def _init_kv_cache(self):
+    return kv_cache.KVCache.from_model_config(self.model.config.decoder_config)
+
+  def _get_extra_args_for_forward(self):
+    return {"called_by_generate": self.forward_called_by_generate}
+
+
+def main(_):
+  if _VERSION.value == "1":
+    checkpoint = _CHECKPOINT[_VERSION.value]
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    checkpoint = kagglehub.model_download(_CHECKPOINT[_VERSION.value])
+    reauthored_checkpoint = checkpoint
+
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = (
+      transformers.PaliGemmaForConditionalGeneration.from_pretrained(checkpoint)
+  )
+
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = paligemma.build_model(
+      reauthored_checkpoint, version=int(_VERSION.value)
+  )
+
+  logging.info("Loading the processor from: %s", checkpoint)
+  # It works only when GemmaTokenizerFast is available. In some environments,
+  # use_fast=False doeesn't work either if the tokenizer cannot load the
+  # sentencepiece model file properly.
+  processor = transformers.AutoProcessor.from_pretrained(checkpoint)
+
+  logging.info("Loading the image from: %s", _IMAGE_URL.value)
+  image = Image.open(requests.get(_IMAGE_URL.value, stream=True).raw)
+  inputs = processor(text=_PROMPTS.value, images=image, return_tensors="pt")
+
+  logging.info("Verifying the reauthored model with model.forward()...")
+  logging.info("Forwarding the original model...")
+  outputs_original = original_model.forward(
+      input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"]
+  )
+  outputs_original = outputs_original.logits
+  logging.info("outputs_original: %s", outputs_original)
+
+  logging.info("Forwarding the reauthored model...")
+  wrapped_reauthored_model = ReauthoredPaliGemmaWrapper(reauthored_model)
+  outputs_reauthored = wrapped_reauthored_model.forward(
+      tokens=inputs["input_ids"],
+      pixel_values=inputs["pixel_values"],
+  )
+  logging.info("outputs_reauthored: %s", outputs_reauthored)
+
+  try:
+    assert torch.allclose(outputs_original, outputs_reauthored, atol=1e-02)
+  except AssertionError as e:
+    logging.error("*** FAILED *** verify with forward()")
+    raise e
+  else:
+    logging.info("*** PASSED *** verify with forward()")
+
+  logging.info("Verifying the reauthored model with model.generate()...")
+  logging.info("Generating answer with the original model...")
+  outputs_original = original_model.generate(
+      **inputs, max_new_tokens=_MAX_NEW_TOKENS.value, do_sample=False
+  )
+  response_original = processor.decode(
+      outputs_original[0], skip_special_tokens=True
+  )
+  logging.info("outputs_from_original_model: [[%s]]", response_original)
+
+  logging.info("Generating answer with the reauthored model...")
+  wrapped_reauthored_model.forward_called_by_generate = True
+  outputs_reauthored = wrapped_reauthored_model.generate(
+      prompts=inputs["input_ids"],
+      pixel_values=inputs["pixel_values"],
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+  )
+  response_reauthored = processor.decode(
+      outputs_reauthored[0], skip_special_tokens=True
+  )
+  logging.info("outputs from reauthored model: [[%s]]", response_reauthored)
+
+  try:
+    assert response_original == response_reauthored
+  except AssertionError as e:
+    logging.error("*** FAILED *** verify with generate()")
+    raise e
+  else:
+    logging.info("*** PASSED *** verify with generate()")
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/paligemma/verify_decoder.py

@@ -0,0 +1,75 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Verifies the reauthored decoder of PaliGemma 3B model."""
+
+import logging
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.paligemma import decoder
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+def main(_):
+  checkpoint = "google/paligemma-3b-mix-224"
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_full_model = (
+      transformers.PaliGemmaForConditionalGeneration.from_pretrained(checkpoint)
+  )
+  original_language_model = original_full_model.eval().language_model
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = decoder.build_decoder(reauthored_checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  # It works only when GemmaTokenizerFast is available. In some environments,
+  # use_fast=False doeesn't work either if the tokenizer cannot load the
+  # sentencepiece model file properly.
+  processor = transformers.AutoProcessor.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_language_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(processor.tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/paligemma/verify_decoder2.py

@@ -0,0 +1,72 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Verifies the reauthored decoder of PaliGemma2 3B model."""
+
+import logging
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.paligemma import decoder2
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import kagglehub
+import transformers
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+def main(_):
+  checkpoint = kagglehub.model_download(
+      "google/paligemma-2/transformers/paligemma2-3b-pt-224"
+  )
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_full_model = (
+      transformers.PaliGemmaForConditionalGeneration.from_pretrained(checkpoint)
+  )
+  original_language_model = original_full_model.eval().language_model
+
+  logging.info("Building the reauthored model from: %s", checkpoint)
+  reauthored_model = decoder2.build_decoder2(checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  # It works only when GemmaTokenizerFast is available. In some environments,
+  # use_fast=False doeesn't work either if the tokenizer cannot load the
+  # sentencepiece model file properly.
+  processor = transformers.AutoProcessor.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_language_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(processor.tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)