Enhance embedding introspection for GenAI config (#321)

Author: apsonawane

examples/gemma4_genai.py

@@ -77,8 +77,10 @@ def build_and_export(
     output_dir: str,
     *,
     dtype: str = "f32",
+    device: str = "cpu",
 ) -> None:
     """Build ONNX model with mobius CLI and write ORT GenAI config."""
+    ep = "cuda" if device == "cuda" else "cpu"
     cmd = [
         sys.executable,
         "-m",
@@ -88,6 +90,8 @@ def build_and_export(
         model_id,
         "--dtype",
         dtype,
+        "--ep",
+        ep,
         "--optimize",
         "--runtime",
         "ort-genai",
@@ -351,7 +355,7 @@ def main() -> None:
 
     # ----- Export-only path -----
     if args.save_to:
-        build_and_export(args.model, args.save_to, dtype=args.dtype)
+        build_and_export(args.model, args.save_to, dtype=args.dtype, device=args.device)
         return
 
     # ----- Resolve model directory -----
@@ -362,7 +366,7 @@ def main() -> None:
         default_dir = os.path.join("output", f"gemma4_{suffix}")
         model_dir = default_dir
         if not os.path.isfile(os.path.join(model_dir, "genai_config.json")):
-            build_and_export(args.model, model_dir, dtype=args.dtype)
+            build_and_export(args.model, model_dir, dtype=args.dtype, device=args.device)
 
     # ----- Inference -----
     print("=" * 60)

src/mobius/integrations/ort_genai/auto_export.py

@@ -148,6 +148,17 @@ def _introspect_inputs(pkg: ModelPackage, key: str) -> dict[str, str] | None:
     return {n: n for n in _graph_input_names(model)}
 
 
+def _introspect_outputs(pkg: ModelPackage, key: str) -> dict[str, str] | None:
+    """Return ``{name: name}`` identity mapping for a sub-model's outputs.
+
+    Returns ``None`` when *key* is absent from *pkg*.
+    """
+    model = pkg.get(key)
+    if model is None:
+        return None
+    return {out.name: out.name for out in model.graph.outputs if out.name is not None}
+
+
 def _copy_tokenizer_files(
     model_id: str,
     output_dir: str,
@@ -713,6 +724,10 @@ def _write_genai_config(
             if embedding_input_mapping is not None:
                 vision_kwargs["embedding_input_names"] = embedding_input_mapping
 
+            embedding_output_mapping = _introspect_outputs(pkg, "embedding")
+            if embedding_output_mapping is not None:
+                vision_kwargs["embedding_output_names"] = embedding_output_mapping
+
             generator.with_vision(image_token_id=image_token_id, **vision_kwargs)
 
     if has_speech:

src/mobius/integrations/ort_genai/genai_config.py

@@ -266,6 +266,7 @@ def with_vision(
         input_names: dict[str, str] | None = None,
         output_names: dict[str, str] | None = None,
         embedding_input_names: dict[str, str] | None = None,
+        embedding_output_names: dict[str, str] | None = None,
         vision_start_token_id: int | None = None,
         video_token_id: int | None = None,
     ) -> GenaiConfigGenerator:
@@ -288,6 +289,8 @@ def with_vision(
                 mapping.  When provided (e.g. from ONNX graph
                 introspection), used directly.  Defaults to
                 input_ids + image_features.
+            embedding_output_names: Override embedding model output name
+                mapping. Defaults to inputs_embeds.
             vision_start_token_id: Token ID for ``<|vision_start|>``.
             video_token_id: Token ID for video placeholders.
 
@@ -321,7 +324,9 @@ def with_vision(
         self._embedding = {
             "filename": embedding_filename,
             "inputs": embedding_input_names,
-            "outputs": {
+            "outputs": embedding_output_names
+            if embedding_output_names is not None
+            else {
                 "inputs_embeds": "inputs_embeds",
             },
             "session_options": _make_session_options(self.ep),