No public description

PiperOrigin-RevId: 865674419

Author: sirakiin

litert_torch/generative/export_hf/core/export_lib.py

@@ -80,6 +80,7 @@ def load_model(
     model_path: str,
     trust_remote_code: bool = False,
     auto_model_override: str | None = None,
+    task: str = 'text_generation',
 ):
   """Loads model from checkpoint."""
 
@@ -90,7 +91,12 @@ def load_model(
   )
   config._attn_implementation = 'lrt_transposed_attention'  # pylint: disable=protected-access
 
-  auto_model_cls = transformers.AutoModelForCausalLM
+  if task == 'text_generation':
+    auto_model_cls = transformers.AutoModelForCausalLM
+  elif task == 'image_text_to_text':
+    auto_model_cls = transformers.AutoModelForImageTextToText
+  else:
+    raise ValueError(f'Unsupported task: {task}')
   if auto_model_override is not None:
     auto_model_cls = transformers.__dict__[auto_model_override]
 
@@ -101,14 +107,16 @@ def load_model(
       trust_remote_code=trust_remote_code,
   )
 
-  model.generation_config.cache_implementation = 'static'
-  model.generation_config.do_sample = False
+  if task == 'text_generation':
+    model.generation_config.cache_implementation = 'static'
+    model.generation_config.do_sample = False
 
   text_model_config = config
   if hasattr(config, 'text_config'):
     text_model_config = config.text_config
 
-  verify_model_compatibility(model, config, text_model_config)
+  if task == 'text_generation':
+    verify_model_compatibility(model, config, text_model_config)
 
   # TODO(weiyiw): Refactor into a separate function.
   tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)

litert_torch/generative/export_hf/core/litert_lm_builder.py

@@ -118,7 +118,8 @@ def build_llm_metadata(
       if isinstance(gen_config.eos_token_id, int):
         stop_tokens.add(gen_config.eos_token_id)
       elif isinstance(gen_config.eos_token_id, list):
-        stop_tokens.update(gen_config.eos_token_id)
+        for token_id in gen_config.eos_token_id:
+          stop_tokens.add(token_id)
     elif hasattr(tokenizer, 'eos_token') and tokenizer.eos_token:
       stop_tokens.add(tokenizer.eos_token)
     for stop_token in stop_tokens:

litert_torch/generative/export_hf/core/patches.py

@@ -60,3 +60,32 @@ def decorator(cls):
 transformers.integrations.use_kernel_forward_from_hub = (
     _use_kernel_forward_from_hub
 )
+
+
+# TODO(weiyiw): Find a better way to patch Gemma3RMSNorm.
+class Gemma3RMSNorm(torch.nn.Module):
+  """RMSNorm Layer."""
+
+  def __init__(self, dim: int, eps: float = 1e-6):
+    """RMSNorm Layer."""
+    super().__init__()
+    self.weight = torch.nn.Parameter(torch.ones(dim))
+    self.variance_epsilon = eps
+    self.hidden_size = dim
+
+  def forward(self, hidden_states):
+    return normalization.rms_norm_with_hlfb(
+        hidden_states,
+        self.weight + 1.0,
+        self.variance_epsilon,
+        torch.ones((self.hidden_size,), dtype=torch.float32),
+    )
+
+  def extra_repr(self):
+    return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+from transformers.models.gemma3 import modeling_gemma3
+
+original_gemma3_rms_norm = modeling_gemma3.Gemma3RMSNorm
+modeling_gemma3.Gemma3RMSNorm = Gemma3RMSNorm

litert_torch/generative/export_hf/export.py

@@ -39,6 +39,7 @@ def export(
     # target_accelerator: str | None = None,
     trust_remote_code: bool = False,
     use_jinja_template: bool = False,
+    task: str = 'text_generation',
 ):
   """Exports HuggingFace Transformers model to tflite."""
   # TODO(weiyiw): Use tmp dir for work_dir.
@@ -48,6 +49,7 @@ def export(
       model,
       trust_remote_code=trust_remote_code,
       auto_model_override=auto_model_override,
+      task=task,
   )
   del config  # Unused.
   if split_cache and not externalize_embedder: