More fixes?

Author: larryliu0820

.ci/scripts/test_huggingface_optimum_model.py

@@ -230,18 +230,13 @@ def test_llm_with_image_modality(
         return_tensors="pt",
     )
 
-    import torch
-
-    from executorch.extension.llm.runner import (
-        GenerationConfig,
-        make_image_input,
-        make_token_input,
-        MultimodalRunner,
-    )
+    from executorch.extension.llm.runner import GenerationConfig, MultimodalRunner
 
     runner = MultimodalRunner(f"{model_dir}/model.pte", f"{model_dir}/tokenizer.model")
     generated_text = runner.generate_text_hf(
-        inputs, GenerationConfig(max_new_tokens=128, temperature=0, echo=False), processor.image_token_id
+        inputs,
+        GenerationConfig(max_new_tokens=128, temperature=0, echo=False),
+        processor.image_token_id,
     )
     print(f"\nGenerated text:\n\t{generated_text}")
     # Free memory before loading eager for quality check

extension/llm/runner/__init__.py

@@ -32,9 +32,10 @@
 
 
 import logging
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Callable, List, Optional, Union
 
 import torch
+from transformers.feature_extraction_utils import BatchFeature
 
 
 def _find_image_token_runs(
@@ -65,13 +66,13 @@ def _find_image_token_runs(
 
 
 def _hf_to_multimodal_inputs(  # noqa: C901
-    inputs: Dict[str, Any], image_token_id: Optional[int] = None
+    inputs: BatchFeature, image_token_id: Optional[int] = None
 ) -> List[MultimodalInput]:
     """Convert a HuggingFace AutoProcessor dict to ExecuTorch MultimodalInputs.
     Currently only support 1 image inside the input.
 
     Args:
-      - inputs: A dictionary containing the input data.
+      - inputs: A BatchFeature containing the input data.
       - image_token_id: The token ID for the image, if present.
 
     `inputs` expected keys:
@@ -168,38 +169,50 @@ def _hf_to_multimodal_inputs(  # noqa: C901
 
 def generate_hf(
     runner: MultimodalRunner,
-    inputs: Union[Dict[str, Any], List[MultimodalInput]],
+    inputs: Union[BatchFeature, List[MultimodalInput]],
     config: GenerationConfig,
     image_token_id: Optional[int] = None,
     token_callback: Optional[Callable[[str], None]] = None,
     stats_callback: Optional[Callable[[Stats], None]] = None,
 ) -> None:
-    """Generate using an HF dict by converting to multimodal inputs internally, or using a list of MultimodalInput."""
-    if isinstance(inputs, dict):
+    """Generate using an BatchFeature by converting to multimodal inputs internally, or using a list of MultimodalInput."""
+    if isinstance(inputs, BatchFeature):
         logging.info(
-            "Input is a dict, assuming it's coming from HF AutoProcessor.apply_chat_template(). Converting to multimodal inputs."
+            "Input is a BatchFeature, assuming it's coming from HF AutoProcessor.apply_chat_template(). Converting to multimodal inputs."
         )
         converted = _hf_to_multimodal_inputs(inputs, image_token_id=image_token_id)
-    else:
+    elif isinstance(inputs, list) and all(
+        isinstance(i, MultimodalInput) for i in inputs
+    ):
         converted = inputs
+    else:
+        raise RuntimeError(
+            "inputs must be either a BatchFeature (from HF AutoProcessor) or a list of MultimodalInput"
+        )
 
     runner.generate(converted, config, token_callback, stats_callback)
 
 
 def generate_text_hf(
     runner: MultimodalRunner,
-    inputs: Union[Dict[str, Any], List[MultimodalInput]],
+    inputs: Union[BatchFeature, List[MultimodalInput]],
     config: GenerationConfig,
     image_token_id: Optional[int] = None,
 ) -> str:
-    """Generate using an HF dict by converting to multimodal inputs internally, or using a list of MultimodalInput."""
-    if isinstance(inputs, dict):
+    """Generate using an BatchFeature by converting to multimodal inputs internally, or using a list of MultimodalInput."""
+    if isinstance(inputs, BatchFeature):
         logging.info(
-            "Input is a dict, assuming it's coming from HF AutoProcessor.apply_chat_template(). Converting to multimodal inputs."
+            "Input is a BatchFeature, assuming it's coming from HF AutoProcessor.apply_chat_template(). Converting to multimodal inputs."
         )
         converted = _hf_to_multimodal_inputs(inputs, image_token_id=image_token_id)
-    else:
+    elif isinstance(inputs, list) and all(
+        isinstance(i, MultimodalInput) for i in inputs
+    ):
         converted = inputs
+    else:
+        raise RuntimeError(
+            "inputs must be either a BatchFeature (from HF AutoProcessor) or a list of MultimodalInput"
+        )
 
     return runner.generate_text(converted, config)
 

extension/llm/runner/_llm_runner.pyi

@@ -4,7 +4,7 @@ Type stubs for _llm_runner module.
 This file provides type annotations for the ExecuTorch LLM Runner Python bindings.
 """
 
-from typing import Callable, List, Optional, Union
+from typing import Callable, List, Optional, Union, overload
 
 import numpy as np
 import torch
@@ -134,14 +134,17 @@ class Stats:
 class Image:
     """Container for image data."""
 
+    @overload
     def __init__(self) -> None:
         """Initialize an empty Image."""
         ...
 
+    @overload
     def __init__(self, data: List[int], width: int, height: int, channels: int) -> None:
         """Initialize an Image with uint8 data."""
         ...
 
+    @overload
     def __init__(
         self, data: List[float], width: int, height: int, channels: int
     ) -> None:
@@ -198,10 +201,12 @@ class Audio:
     n_frames: int
     """Number of time frames."""
 
+    @overload
     def __init__(self) -> None:
         """Initialize an empty Audio."""
         ...
 
+    @overload
     def __init__(
         self, data: List[int], batch_size: int, n_bins: int, n_frames: int
     ) -> None:
@@ -225,10 +230,12 @@ class RawAudio:
     n_samples: int
     """Number of audio samples."""
 
+    @overload
     def __init__(self) -> None:
         """Initialize an empty RawAudio."""
         ...
 
+    @overload
     def __init__(
         self, data: List[int], batch_size: int, n_channels: int, n_samples: int
     ) -> None:
@@ -240,6 +247,7 @@ class RawAudio:
 class MultimodalInput:
     """Container for multimodal input data (text, image, audio, etc.)."""
 
+    @overload
     def __init__(self, text: str) -> None:
         """
         Create a MultimodalInput with text.
@@ -249,6 +257,7 @@ class MultimodalInput:
         """
         ...
 
+    @overload
     def __init__(self, image: Image) -> None:
         """
         Create a MultimodalInput with an image.
@@ -258,6 +267,7 @@ class MultimodalInput:
         """
         ...
 
+    @overload
     def __init__(self, audio: Audio) -> None:
         """
         Create a MultimodalInput with preprocessed audio.
@@ -267,6 +277,7 @@ class MultimodalInput:
         """
         ...
 
+    @overload
     def __init__(self, raw_audio: RawAudio) -> None:
         """
         Create a MultimodalInput with raw audio.
@@ -347,6 +358,7 @@ class MultimodalRunner:
             RuntimeError: If initialization fails
         """
         ...
+
     def generate(
         self,
         inputs: List[MultimodalInput],
@@ -366,7 +378,7 @@ class MultimodalRunner:
         Raises:
             RuntimeError: If generation fails
         """
-    ...
+        ...
 
     def generate_hf(
         self,
@@ -424,9 +436,11 @@ class MultimodalRunner:
         Raises:
             RuntimeError: If generation fails
         """
-    ...
+        ...
 
-    def generate_text_hf(self, inputs: dict, config: GenerationConfig, image_token_id) -> str:
+    def generate_text_hf(
+        self, inputs: dict, config: GenerationConfig, image_token_id
+    ) -> str:
         """
         Generate text directly from a HuggingFace processor dict and return as string.
 

extension/llm/runner/pybindings.cpp

@@ -173,15 +173,15 @@ PYBIND11_MODULE(_llm_runner, m) {
                       float temperature,
                       int32_t num_bos,
                       int32_t num_eos) {
-    GenerationConfig cfg;
-    cfg.echo = echo;
-    cfg.max_new_tokens = max_new_tokens;
-    cfg.warming = warming;
-    cfg.seq_len = seq_len;
-    cfg.temperature = temperature;
-    cfg.num_bos = num_bos;
-    cfg.num_eos = num_eos;
-    return cfg;
+            GenerationConfig cfg;
+            cfg.echo = echo;
+            cfg.max_new_tokens = max_new_tokens;
+            cfg.warming = warming;
+            cfg.seq_len = seq_len;
+            cfg.temperature = temperature;
+            cfg.num_bos = num_bos;
+            cfg.num_eos = num_eos;
+            return cfg;
           }),
           py::arg("echo") = true,
           py::arg("max_new_tokens") = -1,
@@ -204,12 +204,12 @@ PYBIND11_MODULE(_llm_runner, m) {
           py::arg("num_prompt_tokens"),
           "Resolve the maximum number of new tokens to generate based on constraints")
       .def("__repr__", [](const GenerationConfig& config) {
-    return "<GenerationConfig max_new_tokens=" +
-        std::to_string(config.max_new_tokens) +
-        " seq_len=" + std::to_string(config.seq_len) +
-        " temperature=" + std::to_string(config.temperature) +
-        " echo=" + (config.echo ? "True" : "False") +
-        " warming=" + (config.warming ? "True" : "False") + ">";
+        return "<GenerationConfig max_new_tokens=" +
+            std::to_string(config.max_new_tokens) +
+            " seq_len=" + std::to_string(config.seq_len) +
+            " temperature=" + std::to_string(config.temperature) +
+            " echo=" + (config.echo ? "True" : "False") +
+            " warming=" + (config.warming ? "True" : "False") + ">";
       });
 
   // Bind Stats
@@ -365,10 +365,10 @@ PYBIND11_MODULE(_llm_runner, m) {
           py::init<const std::string&>(),
           py::arg("text"),
           "Create a MultimodalInput with text")
-    .def(
-      py::init<const std::vector<uint64_t>&>(),
-      py::arg("tokens"),
-      "Create a MultimodalInput with pre-tokenized tokens (List[int])")
+      .def(
+          py::init<const std::vector<uint64_t>&>(),
+          py::arg("tokens"),
+          "Create a MultimodalInput with pre-tokenized tokens (List[int])")
       .def(
           py::init<const std::vector<uint64_t>&>(),
           py::arg("tokens"),