diff --git a/src/qai_hub_models/models/vietocr/README.md b/src/qai_hub_models/models/vietocr/README.md new file mode 100644 index 000000000..6a4b55ead --- /dev/null +++ b/src/qai_hub_models/models/vietocr/README.md @@ -0,0 +1,84 @@ +# [VietOCR: Vietnamese text recognition with a vgg19_bn CNN backbone and transformer decoder](https://aihub.qualcomm.com/models/vietocr) + +VietOCR is a text recognition model for Vietnamese, covering the full precomposed tone character set. It pairs a vgg19_bn CNN backbone with a transformer sequence decoder. This contribution exports the recognition CNN backbone for on-device deployment. + +This is based on the implementation of VietOCR found [here](https://github.com/pbcquoc/vietocr). +This repository contains scripts for optimized on-device export suitable to run on Qualcomm® devices. More details on model performance across various devices, can be found [here](https://aihub.qualcomm.com/models/vietocr). + +Qualcomm AI Hub Models uses [Qualcomm AI Hub Workbench](https://workbench.aihub.qualcomm.com) to compile, profile, and evaluate this model. [Sign up](https://myaccount.qualcomm.com/signup) to run these models on a hosted Qualcomm® device. + +## Setup +### 1. Install the package +Install the package via pip: +```bash +# NOTE: 3.10 <= PYTHON_VERSION < 3.14 is supported. +pip install "qai-hub-models[vietocr]" +``` + +### 2. Configure Qualcomm® AI Hub Workbench +Sign-in to [Qualcomm® AI Hub Workbench](https://workbench.aihub.qualcomm.com/) with your +Qualcomm® ID. Once signed in navigate to `Account -> Settings -> API Token`. + +With this API token, you can configure your client to run models on the cloud +hosted devices. +```bash +qai-hub configure --api_token API_TOKEN +``` +Navigate to [docs](https://workbench.aihub.qualcomm.com/docs/) for more information. + +## Run CLI Demo +Run the following simple CLI demo to verify the model is working end to end: + +```bash +python -m qai_hub_models.models.vietocr.demo +``` +More details on the CLI tool can be found with the `--help` option. See +[demo.py](demo.py) for sample usage of the model including pre/post processing +scripts. Please refer to our [general instructions on using +models](../../../#getting-started) for more usage instructions. + +## Export for on-device deployment +To run the model on Qualcomm® devices, you must export the model for use with an edge runtime such as +TensorFlow Lite, ONNX Runtime, or Qualcomm AI Engine Direct. Use the following command to export the model: +```bash +python -m qai_hub_models.models.vietocr.export +``` +Additional options are documented with the `--help` option. + +## Scope of this contribution + +Vietnamese text recognition is currently absent from the AI Hub catalog. VietOCR is a +widely used, Apache-2.0, PyTorch-native Vietnamese recognizer, which makes it a good fit +for the standard trace-and-compile export path. This contribution covers the recognition +**CNN backbone** (vgg19_bn); the transformer decoder is a planned follow-up component. + +## Performance + +Measured on a Samsung Galaxy S25 Ultra (Snapdragon 8 Elite) via Qualcomm AI Hub, +float precision, all layers on the Hexagon NPU (no CPU fallback): + +| Component | On-device latency | NPU layer coverage | +|-----------|-------------------|--------------------| +| vgg19_bn backbone | 4.48 ms | 26 / 26 (100%) | + +An end-to-end recognition accuracy number is not reported here because this contribution +covers only the CNN backbone; accuracy will be reported alongside the decoder follow-up. + +## Engineering note + +The original backbone tail uses `permute(-1, 0, 1)` and `transpose(-1, -2).flatten(2)`. +Negative permutation axes and the implied dynamic reshape do not export to a static +on-device graph. The tail is rebuilt with equivalent static, positive-axis operations +(`transpose(2, 3)`, `permute(2, 0, 1)`), preserving semantics while producing a fully +static graph. See `model.py` and `test.py`. + +## License +* The license for the original implementation of VietOCR can be found + [here](https://github.com/pbcquoc/vietocr/blob/master/LICENSE). + +## References +* [Source Model Implementation](https://github.com/pbcquoc/vietocr) + +## Community +* Join [our AI Hub Slack community](https://aihub.qualcomm.com/community/slack) to collaborate, post questions and learn more about on-device AI. +* For questions or feedback please [reach out to us](mailto:ai-hub-support@qti.qualcomm.com). diff --git a/src/qai_hub_models/models/vietocr/__init__.py b/src/qai_hub_models/models/vietocr/__init__.py new file mode 100644 index 000000000..fa3ec652e --- /dev/null +++ b/src/qai_hub_models/models/vietocr/__init__.py @@ -0,0 +1,10 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- + +from .app import VietOCRApp as App +from .model import MODEL_ID +from .model import VietOCR as Model + +__all__ = ["MODEL_ID", "App", "Model"] diff --git a/src/qai_hub_models/models/vietocr/app.py b/src/qai_hub_models/models/vietocr/app.py new file mode 100644 index 000000000..c04a6e34c --- /dev/null +++ b/src/qai_hub_models/models/vietocr/app.py @@ -0,0 +1,140 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from qai_hub_models.models.vietocr.model import IMAGE_HEIGHT, IMAGE_WIDTH + + +class VietOCRApp: + """ + Light-weight "app code" for running the VietOCR recognition CNN backbone. + + For a given text-line image, the app will: + * convert to RGB and resize to the backbone input shape + * normalize pixel values to [0, 1] + * run the backbone to produce a per-column feature sequence + """ + + def __init__( + self, + backbone: Callable[[torch.Tensor], torch.Tensor], + image_shape: tuple[int, int] = (IMAGE_HEIGHT, IMAGE_WIDTH), + ) -> None: + self.backbone = backbone + self.image_shape = image_shape + + def preprocess_image(self, image: Image.Image) -> torch.Tensor: + """Resize and normalize a PIL image into a backbone input tensor.""" + height, width = self.image_shape + resized = image.convert("RGB").resize((width, height)) + arr = np.asarray(resized, dtype=np.float32) / 255.0 + chw = arr.transpose(2, 0, 1)[np.newaxis] + return torch.from_numpy(np.ascontiguousarray(chw)) + + def predict(self, *args: Any, **kwargs: Any) -> np.ndarray: + return self.predict_features_from_image(*args, **kwargs) + + def predict_features_from_image( + self, pixel_values_or_image: np.ndarray | torch.Tensor | Image.Image + ) -> np.ndarray: + """ + Produce the per-column feature sequence for a text-line image. + + Parameters + ---------- + pixel_values_or_image + Input PIL image (before pre-processing), or a tensor / array + already shaped [batch, 3, H, W] with values in [0, 1]. + + Returns + ------- + features : np.ndarray + Per-column feature sequence. Shape [W', batch, 256]. + """ + pixel_values = self._to_input_tensor(pixel_values_or_image) + return np.asarray(self.backbone(pixel_values)) + + def _to_input_tensor( + self, pixel_values_or_image: np.ndarray | torch.Tensor | Image.Image + ) -> torch.Tensor: + if isinstance(pixel_values_or_image, Image.Image): + return self.preprocess_image(pixel_values_or_image) + if isinstance(pixel_values_or_image, np.ndarray): + return torch.from_numpy(pixel_values_or_image) + return pixel_values_or_image + + # ------------------------------------------------------------------ + # End-to-end recognition. + # + # The AI Hub model exported here is the *CNN backbone* (the first of the + # two VietOCR recognizer components). To turn the backbone's per-column + # feature sequence into recognized text, those features are fed through the + # *second* component -- the Transformer seq2seq encoder/decoder plus the + # output vocabulary -- taken from the installed ``vietocr`` package. The + # combination demonstrates full end-to-end Vietnamese text recognition + # while keeping the exported on-device artifact limited to the backbone. + # ------------------------------------------------------------------ + def recognize_text( + self, + image: Image.Image, + transformer: torch.nn.Module, + vocab: Any, + max_seq_length: int = 128, + sos_token: int = 1, + eos_token: int = 2, + ) -> str: + """ + Recognize the Vietnamese text in a text-line image, end to end. + + The backbone (this AI Hub model) produces the feature sequence; the + ``transformer`` (VietOCR's Transformer seq2seq head) greedily decodes + that sequence into token ids, which ``vocab`` maps back to characters. + + Parameters + ---------- + image + A text-line PIL image. + transformer + VietOCR's ``LanguageTransformer`` head (provides + ``forward_encoder`` / ``forward_decoder``). + vocab + VietOCR's ``Vocab`` (provides ``decode``). + max_seq_length, sos_token, eos_token + Greedy-decode controls matching VietOCR's defaults. + + Returns + ------- + text : str + The recognized Vietnamese text. + """ + pixel_values = self._to_input_tensor(image) + + with torch.no_grad(): + # Component 1: exported CNN backbone -> per-column features. + src = self.backbone(pixel_values) + if isinstance(src, np.ndarray): + src = torch.from_numpy(src) + + # Component 2: Transformer seq2seq decode of the features. + memory = transformer.forward_encoder(src) + + translated = [sos_token] + for _ in range(max_seq_length): + tgt = torch.LongTensor([translated]).transpose(0, 1) + output, memory = transformer.forward_decoder(tgt, memory) + next_token = int(output[:, -1, :].argmax(dim=-1).item()) + translated.append(next_token) + if next_token == eos_token: + break + + return vocab.decode(translated) diff --git a/src/qai_hub_models/models/vietocr/code-gen.yaml b/src/qai_hub_models/models/vietocr/code-gen.yaml new file mode 100644 index 000000000..2bc25144e --- /dev/null +++ b/src/qai_hub_models/models/vietocr/code-gen.yaml @@ -0,0 +1,2 @@ +supported_precisions: +- float diff --git a/src/qai_hub_models/models/vietocr/conftest.py b/src/qai_hub_models/models/vietocr/conftest.py new file mode 100644 index 000000000..0b27816ca --- /dev/null +++ b/src/qai_hub_models/models/vietocr/conftest.py @@ -0,0 +1,38 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- +# THIS FILE WAS AUTO-GENERATED. DO NOT EDIT MANUALLY. + +import gc +import warnings + +import pytest +import torch.jit._trace + +from qai_hub_models.models.vietocr import Model +from qai_hub_models.scorecard.utils.testing import make_cached_from_pretrained_fixture + + +def pytest_configure(config: pytest.Config) -> None: + # pytest is unable to figure out how to silence several PyTorch warning types from pyproject.toml settings, + # so we apply a manual warning filter here instead. + warnings.filterwarnings(action="ignore", category=torch.jit._trace.TracerWarning) + warnings.filterwarnings(action="ignore", category=UserWarning, module="torch.*") + warnings.filterwarnings(action="ignore", category=FutureWarning, module="torch.*") + warnings.filterwarnings( + action="ignore", category=DeprecationWarning, module="torch.*" + ) + + +# Instantiate the model only once for all tests. +# Mock from_pretrained to always return the initialized model. +# This speeds up tests and limits memory leaks. +cached_from_pretrained = make_cached_from_pretrained_fixture( + Model, skip_clone_repo=True +) + + +@pytest.fixture(scope="module", autouse=True) +def ensure_gc() -> None: + gc.collect() diff --git a/src/qai_hub_models/models/vietocr/demo.py b/src/qai_hub_models/models/vietocr/demo.py new file mode 100644 index 000000000..dd5e42bca --- /dev/null +++ b/src/qai_hub_models/models/vietocr/demo.py @@ -0,0 +1,97 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- + +from __future__ import annotations + +from typing import Any + +from qai_hub_models.models.vietocr.app import VietOCRApp +from qai_hub_models.models.vietocr.model import ( + MODEL_ASSET_VERSION, + MODEL_ID, + VIETOCR_CONFIG_NAME, + VietOCR, +) +from qai_hub_models.utils.args import get_model_cli_parser, model_from_cli_args +from qai_hub_models.utils.asset_loaders import CachedWebModelAsset, load_image + +DEFAULT_SAMPLE_IMAGE = CachedWebModelAsset.from_asset_store( + MODEL_ID, MODEL_ASSET_VERSION, "sample_text.jpg" +) + + +def _load_transformer_head_and_vocab() -> tuple[Any, Any, Any]: + """Load VietOCR's Transformer seq2seq head and output vocab. + + The AI Hub model in this directory is only the CNN backbone (the first of + VietOCR's two recognizer components). The Transformer seq2seq head plus its + output vocabulary -- the second component -- are loaded here from the + installed ``vietocr`` package so the demo can decode the backbone features + into recognized text end to end. + """ + import torch + from vietocr.model.transformerocr import VietOCR as _VietOCR + from vietocr.model.vocab import Vocab + from vietocr.tool.config import Cfg + from vietocr.tool.utils import download_weights + + cfg = Cfg.load_config_from_name(VIETOCR_CONFIG_NAME) + cfg["device"] = "cpu" + vocab = Vocab(cfg["vocab"]) + full_model = _VietOCR( + len(vocab), + cfg["backbone"], + cfg["cnn"], + cfg["transformer"], + cfg["seq_modeling"], + ).eval() + weights = download_weights(cfg["pretrain"]) + full_model.load_state_dict(torch.load(weights, map_location="cpu")) + return full_model, full_model.transformer, vocab + + +def main(is_test: bool = False) -> None: + import numpy as np + + # Demo parameters + parser = get_model_cli_parser(VietOCR) + parser.add_argument( + "--image", + type=str, + default=DEFAULT_SAMPLE_IMAGE, + help="image file path or URL", + ) + args = parser.parse_args([] if is_test else None) + + # Load the exported CNN backbone (AI Hub model) and wrap it in the app. + app = VietOCRApp(model_from_cli_args(VietOCR, args)) + + # Load image and run the backbone -> per-column feature sequence. + image = load_image(args.image) + features = app.predict_features_from_image(image) + + # Load the second recognizer component (Transformer head + vocab) from the + # vietocr package and decode the backbone features into text end to end. + full_model, transformer, vocab = _load_transformer_head_and_vocab() + text = app.recognize_text(image, transformer, vocab) + + # Sanity check: the exported backbone reproduces the package backbone's + # features on the same input, proving the on-device CNN export is correct. + pixel_values = app.preprocess_image(image) + with __import__("torch").no_grad(): + reference = np.asarray(full_model.cnn(pixel_values)) + backbone_features = app.predict_features_from_image(pixel_values) + features_match = np.allclose(backbone_features, reference, atol=1e-4) + + if not is_test: + print(f"Backbone feature sequence shape: {features.shape}") + print(f"Exported backbone matches package backbone: {features_match}") + print(f"Recognized text: {text}") + else: + assert features_match, "Exported backbone features diverge from VietOCR's." + + +if __name__ == "__main__": + main() diff --git a/src/qai_hub_models/models/vietocr/export.py b/src/qai_hub_models/models/vietocr/export.py new file mode 100644 index 000000000..68a7037a6 --- /dev/null +++ b/src/qai_hub_models/models/vietocr/export.py @@ -0,0 +1,552 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- +# THIS FILE WAS AUTO-GENERATED. DO NOT EDIT MANUALLY. + + +from __future__ import annotations + +import os +import shutil +import tempfile +import warnings +from pathlib import Path +from typing import Any + +import qai_hub as hub + +from qai_hub_models import Precision, TargetRuntime +from qai_hub_models.common import SampleInputsType +from qai_hub_models.configs.model_metadata import ( + ChipsetAttributes, + ModelFileMetadata, + ModelMetadata, + merge_input_metadata, + merge_output_metadata, +) +from qai_hub_models.configs.tool_versions import ToolVersions +from qai_hub_models.models.vietocr import MODEL_ID, Model +from qai_hub_models.utils import quantization as quantization_utils +from qai_hub_models.utils.args import ( + export_parser, + get_export_model_name, + get_input_spec_kwargs, + get_model_kwargs, +) +from qai_hub_models.utils.asset_loaders import ASSET_CONFIG +from qai_hub_models.utils.base_model import BaseModel +from qai_hub_models.utils.compare import torch_inference +from qai_hub_models.utils.export_result import ExportResult +from qai_hub_models.utils.export_without_hub_access import export_without_hub_access +from qai_hub_models.utils.input_spec import InputSpec, to_hub_input_specs +from qai_hub_models.utils.onnx.helpers import download_and_unzip_workbench_onnx_model +from qai_hub_models.utils.path_helpers import get_next_free_path +from qai_hub_models.utils.printing import ( + print_inference_metrics, + print_on_target_demo_cmd, + print_profile_metrics_from_job, + print_tool_versions, +) +from qai_hub_models.utils.qai_hub_helpers import can_access_qualcomm_ai_hub + + +def quantize_model( + precision: Precision, + model: BaseModel, + model_name: str, + onnx_model: hub.Model, + num_calibration_samples: int | None, + extra_options: str = "", + input_spec: InputSpec | None = None, +) -> hub.client.QuantizeJob: + input_spec = input_spec or model.get_input_spec() + print(f"Quantizing {model_name}.") + if not precision.activations_type or not precision.weights_type: + raise ValueError( + "Quantization is only supported if both weights and activations are quantized." + ) + + calibration_data = quantization_utils.get_calibration_data( + model, input_spec, num_calibration_samples + ) + return hub.submit_quantize_job( + model=onnx_model, + calibration_data=calibration_data, + activations_dtype=precision.activations_type, + weights_dtype=precision.weights_type, + name=model_name, + options=model.get_hub_quantize_options(precision, extra_options), + ) + + +def upload_model( + model: BaseModel, + input_spec: InputSpec | None = None, +) -> hub.Model: + input_spec = input_spec or model.get_input_spec() + with tempfile.TemporaryDirectory() as tmpdir: + return hub.upload_model(str(model.serialize(tmpdir, input_spec))) + + +def compile_model( + model: BaseModel, + model_name: str, + device: hub.Device, + target_runtime: TargetRuntime, + precision: Precision, + source_model: hub.Model, + input_spec: InputSpec | None = None, + extra_options: str = "", +) -> hub.client.CompileJob: + input_spec = input_spec or model.get_input_spec() + + model_compile_options = model.get_hub_compile_options( + target_runtime, precision, extra_options, device + ) + print(f"Optimizing model {model_name} to run on-device") + return hub.submit_compile_job( + model=source_model, + input_specs=to_hub_input_specs(input_spec), + device=device, + name=model_name, + options=model_compile_options, + ) + + +def link_model( + compiled_model: hub.Model, + device: hub.Device, + model_name: str, + model: BaseModel, + target_runtime: TargetRuntime, + extra_options: str = "", +) -> hub.client.LinkJob: + """Link compiled DLC to context binary for AOT.""" + assert target_runtime.is_aot_compiled, ( + f"link_model() requires an AOT runtime, got {target_runtime}" + ) + link_options = model.get_hub_link_options(target_runtime, extra_options) + print(f"Linking {model_name} to context binary") + return hub.submit_link_job( + [compiled_model], + device=device, + name=model_name, + options=link_options, + ) + + +def profile_model( + model_name: str, + device: hub.Device, + options: str, + target_model: hub.Model, +) -> hub.client.ProfileJob: + print(f"Profiling model {model_name} on a hosted device.") + return hub.submit_profile_job( + model=target_model, + device=device, + name=model_name, + options=options, + ) + + +def inference_model( + inputs: SampleInputsType, + model_name: str, + device: hub.Device, + options: str, + target_model: hub.Model, +) -> hub.client.InferenceJob: + print(f"Running inference for {model_name} on a hosted device with example inputs.") + return hub.submit_inference_job( + model=target_model, + inputs=inputs, + device=device, + name=model_name, + options=options, + ) + + +def download_model( + output_dir: os.PathLike | str, + model: BaseModel, + runtime: TargetRuntime, + precision: Precision, + tool_versions: ToolVersions, + target_model: hub.Model, + model_name: str, + zip_assets: bool, + hub_device: hub.Device | None = None, +) -> Path: + output_folder_name = os.path.basename(output_dir) + output_path = get_next_free_path(output_dir) + + with tempfile.TemporaryDirectory() as tmpdir: + dst_path = Path(tmpdir) / output_folder_name + dst_path.mkdir() + + if target_model.model_type == hub.SourceModelType.ONNX: + onnx_result = download_and_unzip_workbench_onnx_model( + target_model, dst_path, model_name + ) + model_file_name = onnx_result.onnx_graph_name + else: + downloaded_path = target_model.download(os.path.join(dst_path, model_name)) + model_file_name = os.path.basename(downloaded_path) + + # Extract and save metadata alongside downloaded model + metadata_path = dst_path / "metadata.json" + file_metadata = ModelFileMetadata.from_hub_model(target_model) + # Merge semantic metadata from get_input_spec() + merge_input_metadata(file_metadata, model.get_input_spec()) + merge_output_metadata(file_metadata, model.get_output_spec()) + model_metadata = ModelMetadata( + model_id=MODEL_ID, + model_name="VietOCR", + runtime=runtime, + precision=precision, + tool_versions=tool_versions, + model_files={model_file_name: file_metadata}, + chipset_attributes=ChipsetAttributes.from_hub_device(hub_device) + if runtime.is_aot_compiled + else None, + ) + + # Dump supplementary files into the model folder + model.write_supplementary_files(dst_path, model_metadata) + + model_metadata.to_json(metadata_path) + if zip_assets: + output_path = Path( + shutil.make_archive( + str(output_path), + "zip", + root_dir=tmpdir, + base_dir=output_folder_name, + ) + ) + else: + shutil.move(dst_path, output_path) + + return output_path + + +def export_model( + device: hub.Device, + precision: Precision = Precision.float, + num_calibration_samples: int | None = None, + quantized_model_id: str | None = None, + skip_compiling: bool = False, + skip_profiling: bool = False, + skip_inferencing: bool = False, + skip_downloading: bool = False, + skip_summary: bool = False, + output_dir: str | None = None, + target_runtime: TargetRuntime = TargetRuntime.TFLITE, + compile_options: str = "", + quantize_options: str = "", + profile_options: str = "", + fetch_static_assets: str | None = None, + zip_assets: bool = False, + **additional_model_kwargs: Any, +) -> ExportResult: + """ + This function executes the following recipe: + + 1. Instantiates a PyTorch model and converts it to a traced TorchScript format + 2. Converts the PyTorch model to ONNX and quantizes the ONNX model. + 3. Compiles the model to an asset that can be run on device + 4. Profiles the model performance on a real device + 5. Inferences the model on sample inputs + 6. Extracts relevant tool (eg. SDK) versions used to compile and profile this model + 7. Downloads the model asset to the local directory + 8. Summarizes the results from profiling and inference + + Each of the last 6 steps can be optionally skipped using the input options. + + Parameters + ---------- + device + Device for which to export the model (e.g., hub.Device("Samsung Galaxy S25")). + Full list of available devices can be found by running `hub.get_devices()`. + precision + The precision to which this model should be quantized. + Quantization is skipped if the precision is float. + num_calibration_samples + The number of calibration data samples + to use for quantization. If not set, uses the default number + specified by the dataset. If model doesn't have a calibration dataset + specified, this must be None. + quantized_model_id + A quantized ONNX hub model id, skips quantizing model. + skip_compiling + If set, skips compiling of model to format that can run on device. + skip_profiling + If set, skips profiling of compiled model on real devices. + skip_inferencing + If set, skips computing on-device outputs from sample data. + skip_downloading + If set, skips downloading of compiled model. + skip_summary + If set, skips waiting for and summarizing results + from profiling and inference. + output_dir + Directory to store generated assets (e.g. compiled model). + Defaults to `/export_assets`. + target_runtime + Which on-device runtime to target. Default is TFLite. + compile_options + Additional options to pass when submitting the compile job. + quantize_options + Additional options to pass when submitting the quantize job. + profile_options + Additional options to pass when submitting the profile job. + fetch_static_assets + If set, known assets are fetched from the given version rather than re-computing them. Can be passed as "latest" or "v". + zip_assets + If set, zip the assets after downloading. + **additional_model_kwargs + Additional optional kwargs used to customize + `model_cls.from_pretrained` and `model.get_input_spec` + + Returns + ------- + ExportResult + * A CompileJob object containing metadata about the compile job submitted to hub (None if compiling skipped). + * An InferenceJob containing metadata about the inference job (None if inferencing skipped). + * A ProfileJob containing metadata about the profile job (None if profiling skipped). + * A QuantizeJob object containing metadata about the quantize job submitted to hub + * The path to the downloaded model folder (or zip), or None if one or more of: skip_downloading is True, fetch_static_assets is set, or AI Hub Workbench is not accessible + """ + model_name = get_export_model_name( + Model, MODEL_ID, precision, additional_model_kwargs + ) + + output_path = Path(output_dir or Path.cwd() / "export_assets") + if fetch_static_assets or not can_access_qualcomm_ai_hub(): + static_model_path = export_without_hub_access( + MODEL_ID, + device, + skip_profiling, + skip_inferencing, + skip_downloading, + skip_summary, + output_path, + target_runtime, + precision, + quantize_options + compile_options + profile_options, + qaihm_version_tag=fetch_static_assets, + ) + return ExportResult(download_path=static_model_path) + + hub_device = hub.get_devices( + name=device.name, attributes=device.attributes, os=device.os + )[-1] + chipset_attr = next( + (attr for attr in hub_device.attributes if "chipset" in attr), None + ) + chipset = chipset_attr.split(":")[-1] if chipset_attr else None + + # 1. Instantiates a PyTorch model and converts it to a traced TorchScript format + model = Model.from_pretrained( + **get_model_kwargs(Model, dict(**additional_model_kwargs, precision=precision)) + ) + input_spec = model.get_input_spec( + **get_input_spec_kwargs(model, additional_model_kwargs) + ) + source_model_to_compile = upload_model(model, input_spec) + + # 2. Converts the PyTorch model to ONNX and quantizes the ONNX model. + quantize_job: hub.client.QuantizeJob | None = None + quantized_model: hub.Model | None = None + if precision != Precision.float: + if quantized_model_id: + quantized_model = hub.get_model(quantized_model_id) + assert quantized_model is not None + else: + onnx_compile_result = compile_model( + model, + model_name, + device, + TargetRuntime.ONNX, + precision, + source_model_to_compile, + input_spec=input_spec, + ) + onnx_model = onnx_compile_result.get_target_model() + assert onnx_model is not None, ( + f"ONNX compile job failed: {onnx_compile_result}" + ) + quantize_job = quantize_model( + precision, + model, + model_name, + onnx_model, + num_calibration_samples, + quantize_options, + input_spec, + ) + if skip_compiling: + return ExportResult(quantize_job=quantize_job) + quantized_model = quantize_job.get_target_model() + assert quantized_model is not None, f"Quantize job failed: {quantize_job}" + + # 3. Compiles the model to an asset that can be run on device + if quantized_model: + source_model_to_compile = quantized_model + compile_result = compile_model( + model, + model_name, + device, + target_runtime, + precision, + source_model_to_compile, + input_spec=input_spec, + extra_options=compile_options, + ) + + link_result: hub.client.LinkJob | None = None + target_model: hub.Model | None + if target_runtime.uses_hub_link: + compiled_model = compile_result.get_target_model() + assert compiled_model is not None, f"Compile job failed: {compile_result}" + link_result = link_model( + compiled_model, + device, + model_name, + model, + target_runtime, + ) + # Extract target models from link jobs for profile/inference + target_model = link_result.get_target_model() + assert target_model is not None, f"Link job failed: {link_result}" + else: + # For JIT runtimes, extract models from compile jobs + target_model = compile_result.get_target_model() + assert target_model is not None, f"Compile job failed: {compile_result}" + + # 4. Profiles the model performance on a real device + profile_result: hub.client.ProfileJob | None = None + if not skip_profiling: + profile_result = profile_model( + model_name, + device, + model.get_hub_profile_options(target_runtime, profile_options), + target_model, + ) + + # 5. Inferences the model on sample inputs + inference_result: hub.client.InferenceJob | None = None + if not skip_inferencing: + inference_result = inference_model( + model.sample_inputs( + input_spec=input_spec, + use_channel_last_format=target_runtime.channel_last_native_execution, + ), + model_name, + device, + model.get_hub_profile_options(target_runtime, profile_options), + target_model, + ) + + # 6. Extracts relevant tool (eg. SDK) versions used to compile and profile this model + tool_versions: ToolVersions | None = None + tool_versions_are_from_device_job = False + if not skip_summary or not skip_downloading: + if profile_result is not None and profile_result.wait(): + tool_versions = ToolVersions.from_job(profile_result) + tool_versions_are_from_device_job = True + elif inference_result is not None and inference_result.wait(): + tool_versions = ToolVersions.from_job(inference_result) + tool_versions_are_from_device_job = True + elif compile_result and compile_result.wait(): + tool_versions = ToolVersions.from_job(compile_result) + + # 7. Downloads the model asset to the local directory + downloaded_model_path: Path | None = None + if not skip_downloading and tool_versions is not None: + model_directory = output_path / ASSET_CONFIG.get_release_asset_name( + MODEL_ID, target_runtime, precision, chipset + ) + downloaded_model_path = download_model( + model_directory, + model, + target_runtime, + precision, + tool_versions, + target_model, + MODEL_ID, + zip_assets, + hub_device=hub_device, + ) + + # 8. Summarizes the results from profiling and inference + if not skip_summary and profile_result is not None: + assert profile_result.wait().success, "Job failed: " + profile_result.url + profile_data: dict[str, Any] = profile_result.download_profile() + print_profile_metrics_from_job(profile_result, profile_data) + + if not skip_summary and inference_result is not None: + sample_inputs = model.sample_inputs(input_spec, use_channel_last_format=False) + torch_out = torch_inference( + model, + sample_inputs, + return_channel_last_output=target_runtime.channel_last_native_execution, + ) + assert inference_result.wait().success, "Job failed: " + inference_result.url + ij_output = inference_result.download_output_data() + assert ij_output is not None + print_inference_metrics( + inference_result, ij_output, torch_out, model.get_output_names() + ) + + if not skip_summary: + print_tool_versions(tool_versions, tool_versions_are_from_device_job) + print_on_target_demo_cmd( + link_result or compile_result, + Path(__file__).parent, + device, + ) + + if downloaded_model_path: + print(f"{model_name} was saved to {downloaded_model_path}\n") + + return ExportResult( + quantize_job=quantize_job, + compile_job=compile_result, + link_job=link_result, + inference_job=inference_result, + profile_job=profile_result, + download_path=downloaded_model_path, + tool_versions=tool_versions, + ) + + +def main() -> None: + warnings.filterwarnings("ignore") + supported_precision_runtimes: dict[Precision, list[TargetRuntime]] = { + Precision.float: [ + TargetRuntime.TFLITE, + TargetRuntime.QNN_DLC, + TargetRuntime.QNN_CONTEXT_BINARY, + TargetRuntime.ONNX, + TargetRuntime.PRECOMPILED_QNN_ONNX, + ], + Precision.w8a8: [ + TargetRuntime.ONNX, + ], + } + + parser = export_parser( + model_cls=Model, + export_fn=export_model, + supported_precision_runtimes=supported_precision_runtimes, + default_export_device="Samsung Galaxy S25 (Family)", + ) + args = parser.parse_args() + export_model(**vars(args)) + + +if __name__ == "__main__": + main() diff --git a/src/qai_hub_models/models/vietocr/info.yaml b/src/qai_hub_models/models/vietocr/info.yaml new file mode 100644 index 000000000..57e31b1c7 --- /dev/null +++ b/src/qai_hub_models/models/vietocr/info.yaml @@ -0,0 +1,28 @@ +name: VietOCR +id: vietocr +status: pending +headline: Vietnamese text recognition with a vgg19_bn CNN backbone and transformer decoder. +domain: Multimodal +description: VietOCR is a text recognition model for Vietnamese, covering the full precomposed tone character set. It pairs a vgg19_bn CNN backbone with a transformer sequence decoder. This contribution exports the recognition CNN backbone for on-device deployment. +use_case: Image To Text +tags: [] +applicable_scenarios: +- Document Management +- Offline OCR +related_models: +- trocr +- easyocr +form_factors: +- Phone +- Tablet +has_static_banner: false +has_animated_banner: false +dataset: [] +technical_details: + Source model: VietOCR vgg_transformer (pbcquoc/vietocr) + Backbone: vgg19_bn CNN with a 1x1 projection conv + Input resolution: 32x128 + NPU latency (Samsung Galaxy S25 Ultra, float, 100% NPU): 4.48 ms +license_type: apache-2.0 +source_repo: https://github.com/pbcquoc/vietocr +license: https://github.com/pbcquoc/vietocr/blob/master/LICENSE diff --git a/src/qai_hub_models/models/vietocr/model.py b/src/qai_hub_models/models/vietocr/model.py new file mode 100644 index 000000000..fb15561e3 --- /dev/null +++ b/src/qai_hub_models/models/vietocr/model.py @@ -0,0 +1,127 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from typing_extensions import Self + +from qai_hub_models.utils.base_model import BaseModel +from qai_hub_models.utils.input_spec import ( + ColorFormat, + ImageMetadata, + InputSpec, + IoType, + TensorSpec, +) + +MODEL_ID = __name__.split(".")[-2] +MODEL_ASSET_VERSION = 1 + +# VietOCR's default vgg_transformer recognizer. The recognition CNN backbone +# (vgg19_bn) is the convolution-heavy front end of the model. +VIETOCR_CONFIG_NAME = "vgg_transformer" + +# Recognition is performed on fixed-height text-line crops; a representative +# fixed width is pinned for on-device export. +IMAGE_HEIGHT = 32 +IMAGE_WIDTH = 128 + + +def _load_vietocr_cnn() -> torch.nn.Module: + """Load the pretrained VietOCR vgg_transformer recognizer's CNN backbone.""" + from vietocr.model.transformerocr import VietOCR as _VietOCR + from vietocr.model.vocab import Vocab + from vietocr.tool.config import Cfg + + cfg = Cfg.load_config_from_name(VIETOCR_CONFIG_NAME) + cfg["device"] = "cpu" + vocab = Vocab(cfg["vocab"]) + model = _VietOCR( + len(vocab), + cfg["backbone"], + cfg["cnn"], + cfg["transformer"], + cfg["seq_modeling"], + ).eval() + + from vietocr.tool.utils import download_weights + + weights = download_weights(cfg["pretrain"]) + model.load_state_dict(torch.load(weights, map_location="cpu")) + + # The CNN backbone is wrapped in `.model` inside VietOCR's CNN module. + return model.cnn.model if hasattr(model.cnn, "model") else model.cnn + + +class VietOCR(BaseModel): + """VietOCR recognition CNN backbone (vgg19_bn) for Vietnamese text. + + VietOCR recognizes Vietnamese text (a character set that includes the full + precomposed tone vowels, e.g. ế ồ ự ấ ợ). This component is the CNN backbone + that maps a text-line crop to a per-column feature sequence consumed by the + downstream transformer recognizer. + + The original backbone tail applies `permute(-1, 0, 1)` and + `transpose(-1, -2).flatten(2)`. Negative permutation axes and the implied + dynamic reshape do not export cleanly to a static on-device graph, so the + tail is rebuilt here with equivalent static, positive-axis operations. + """ + + def __init__(self, vgg: torch.nn.Module) -> None: + super().__init__() + self.features = vgg.features + self.last_conv_1x1 = vgg.last_conv_1x1 + + @classmethod + def from_pretrained(cls) -> Self: + return cls(_load_vietocr_cnn()) + + def forward(self, image: torch.Tensor) -> torch.Tensor: + """ + Run the CNN backbone on a text-line crop. + + Parameters + ---------- + image + Pixel values pre-processed for backbone consumption. + Range: float[0, 1] + 3-channel Color Space: RGB + + Returns + ------- + features : torch.Tensor + Per-column feature sequence. Shape [W', batch, 256]. + """ + x = self.features(image) # [B, 512, 1, W'] + x = self.last_conv_1x1(x) # [B, 256, 1, W'] + x = x.transpose(2, 3) # [B, 256, W', 1] (static positive axes) + x = x.flatten(2) # [B, 256, W'] + # [W', B, 256] (static positive axes) + return x.permute(2, 0, 1) + + def get_input_spec( + self, + batch_size: int = 1, + height: int = IMAGE_HEIGHT, + width: int = IMAGE_WIDTH, + ) -> InputSpec: + return { + "image": TensorSpec( + shape=(batch_size, 3, height, width), + dtype="float32", + io_type=IoType.IMAGE, + value_range=(0.0, 1.0), + image_metadata=ImageMetadata( + color_format=ColorFormat.RGB, + ), + ), + } + + def get_output_names(self) -> list[str]: + return ["features"] + + def get_channel_last_inputs(self) -> list[str]: + return ["image"] diff --git a/src/qai_hub_models/models/vietocr/requirements.txt b/src/qai_hub_models/models/vietocr/requirements.txt new file mode 100644 index 000000000..391753121 --- /dev/null +++ b/src/qai_hub_models/models/vietocr/requirements.txt @@ -0,0 +1 @@ +vietocr==0.3.13 diff --git a/src/qai_hub_models/models/vietocr/test.py b/src/qai_hub_models/models/vietocr/test.py new file mode 100644 index 000000000..fd5159d89 --- /dev/null +++ b/src/qai_hub_models/models/vietocr/test.py @@ -0,0 +1,63 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- + +import numpy as np +import torch + +from qai_hub_models.models.vietocr.demo import main as demo_main +from qai_hub_models.models.vietocr.model import ( + IMAGE_HEIGHT, + IMAGE_WIDTH, + VietOCR, +) +from qai_hub_models.scorecard.utils.testing import skip_clone_repo_check + + +@skip_clone_repo_check +def test_task() -> None: + model = VietOCR.from_pretrained() + spec = model.get_input_spec() + assert spec["image"][0] == (1, 3, IMAGE_HEIGHT, IMAGE_WIDTH) + + image = torch.rand(*spec["image"][0]) + features = np.asarray(model(image)) + + # Output is a per-column feature sequence: [W', batch, 256]. + assert features.ndim == 3 + assert features.shape[1] == 1 + assert features.shape[2] == 256 + + +@skip_clone_repo_check +def test_static_export_graph() -> None: + """The rebuilt backbone tail must export to a static ONNX graph with no + negative-axis Transpose (the reason the original tail did not compile). + """ + import onnx + + model = VietOCR.from_pretrained() + image = torch.rand(*model.get_input_spec()["image"][0]) + path = "/tmp/vietocr_test.onnx" + torch.onnx.export( + model, + image, + path, + input_names=["image"], + output_names=["features"], + opset_version=17, + do_constant_folding=True, + ) + onnx_model = onnx.load(path) + onnx.checker.check_model(onnx_model) + onnx.shape_inference.infer_shapes(onnx_model, check_type=True, strict_mode=True) + for node in onnx_model.graph.node: + if node.op_type == "Transpose": + for attr in node.attribute: + if attr.name == "perm": + assert all(p >= 0 for p in attr.ints) + + +def test_demo() -> None: + demo_main(is_test=True) diff --git a/src/qai_hub_models/models/vietocr/test_generated.py b/src/qai_hub_models/models/vietocr/test_generated.py new file mode 100644 index 000000000..2e64eb5ac --- /dev/null +++ b/src/qai_hub_models/models/vietocr/test_generated.py @@ -0,0 +1,443 @@ +# --------------------------------------------------------------------- +# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# --------------------------------------------------------------------- +# THIS FILE WAS AUTO-GENERATED. DO NOT EDIT MANUALLY. + +from __future__ import annotations + +import os +from collections.abc import Generator +from pathlib import Path + +import numpy as np +import pytest +import qai_hub as hub +import torch + +from qai_hub_models import Precision, TargetRuntime +from qai_hub_models.models.vietocr import MODEL_ID, Model +from qai_hub_models.models.vietocr.export import ( + compile_model, + export_model, + inference_model, + link_model, + profile_model, + quantize_model, + upload_model, +) +from qai_hub_models.scorecard import ( + ScorecardCompilePath, + ScorecardDevice, + ScorecardProfilePath, +) +from qai_hub_models.scorecard.errors import CachedScorecardJobError +from qai_hub_models.scorecard.execution_helpers import ( + get_compile_parameterized_pytest_config, + get_evaluation_parameterized_pytest_config, + get_export_parameterized_pytest_config, + get_link_parameterized_pytest_config, + get_profile_parameterized_pytest_config, + get_quantize_parameterized_pytest_config, + needs_pre_quantize_compile, + pytest_device_idfn, +) +from qai_hub_models.scorecard.utils.testing import skip_invalid_runtime_device +from qai_hub_models.scorecard.utils.testing_export_eval import ( + accuracy_on_dataset_via_evaluate_and_export, + accuracy_on_sample_inputs_via_export, + compile_via_export, + export_test_e2e, + inference_via_export, + link_via_export, + on_device_inference_for_accuracy_validation, + pre_quantize_compile_via_export, + profile_via_export, + quantize_via_export, + sim_accuracy_on_dataset, + split_and_group_accuracy_validation_output_batches, + torch_accuracy_on_dataset, + torch_inference_for_accuracy_validation, + torch_inference_for_accuracy_validation_outputs, +) +from qai_hub_models.utils.args import get_model_kwargs +from qai_hub_models.utils.input_spec import InputSpec +from qai_hub_models.utils.validation import perform_runtime_model_validation + +# All runtime + precision pairs that are enabled for testing and are compatibile with this model. +# NOTE: +# Certain supported pairs may be excluded from this list if they are not enabled for testing. +# For example, models that allow JIT (on-device) compile will not test AOT runtimes; we assume that if it works on JIT it will work on AOT. +ENABLED_PRECISION_RUNTIMES: dict[Precision, list[TargetRuntime]] = { + Precision.float: [ + TargetRuntime.TFLITE, + TargetRuntime.QNN_DLC, + TargetRuntime.ONNX, + ], + Precision.w8a8: [ + TargetRuntime.TFLITE, + TargetRuntime.QNN_DLC, + TargetRuntime.ONNX, + ], +} + + +# All runtime + precision pairs that are enabled for testing and have no known failure reasons. +# NOTE: +# Certain supported pairs may be excluded from this list if they are not enabled for testing. +# For example, models that allow JIT (on-device) compile will not test AOT runtimes; we assume that if it works on JIT it will work on AOT. +PASSING_PRECISION_RUNTIMES: dict[Precision, list[TargetRuntime]] = { + Precision.float: [ + TargetRuntime.TFLITE, + TargetRuntime.QNN_DLC, + TargetRuntime.ONNX, + ], + Precision.w8a8: [ + TargetRuntime.ONNX, + ], +} + + +EVAL_DEVICE = ScorecardDevice.get("Samsung Galaxy S25 (Family)") +HAS_EVAL_DATASET = len(Model.get_eval_dataset_classes()) > 0 + + +@pytest.mark.compile +def test_runtime_model_validation() -> None: + perform_runtime_model_validation(Model, MODEL_ID) + + +@pytest.mark.pre_quantize_compile +@pytest.mark.skipif( + not needs_pre_quantize_compile( + MODEL_ID, ENABLED_PRECISION_RUNTIMES, PASSING_PRECISION_RUNTIMES + ), + reason="Model does not require pre-quantize compile step", +) +def test_pre_quantize_compile() -> None: + pre_quantize_compile_via_export( + compile_model, + MODEL_ID, + Model.from_pretrained(), + upload_model, + ) + + +@pytest.mark.parametrize( + "precision", + get_quantize_parameterized_pytest_config( + MODEL_ID, ENABLED_PRECISION_RUNTIMES, PASSING_PRECISION_RUNTIMES + ), + ids=pytest_device_idfn, +) +@pytest.mark.quantize +def test_quantize(precision: Precision) -> None: + try: + quantize_via_export( + quantize_model, + MODEL_ID, + Model.from_pretrained(), + precision, + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +@pytest.mark.parametrize( + ("precision", "scorecard_path", "device"), + get_compile_parameterized_pytest_config( + MODEL_ID, ENABLED_PRECISION_RUNTIMES, PASSING_PRECISION_RUNTIMES + ), + ids=pytest_device_idfn, +) +@pytest.mark.compile +def test_compile( + precision: Precision, scorecard_path: ScorecardCompilePath, device: ScorecardDevice +) -> None: + skip_invalid_runtime_device(Model, scorecard_path.runtime, device) + try: + compile_via_export( + compile_model, + MODEL_ID, + Model.from_pretrained(), + precision, + scorecard_path, + device, + upload_model=upload_model, + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +@pytest.mark.parametrize( + ("precision", "scorecard_path", "device"), + get_link_parameterized_pytest_config( + MODEL_ID, ENABLED_PRECISION_RUNTIMES, PASSING_PRECISION_RUNTIMES + ), + ids=pytest_device_idfn, +) +@pytest.mark.link +def test_link( + precision: Precision, scorecard_path: ScorecardCompilePath, device: ScorecardDevice +) -> None: + skip_invalid_runtime_device(Model, scorecard_path.runtime, device) + try: + link_via_export( + link_model, + MODEL_ID, + Model.from_pretrained(), + precision, + scorecard_path, + device, + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +@pytest.mark.parametrize( + ("precision", "scorecard_path", "device"), + get_profile_parameterized_pytest_config( + MODEL_ID, ENABLED_PRECISION_RUNTIMES, PASSING_PRECISION_RUNTIMES + ), + ids=pytest_device_idfn, +) +@pytest.mark.profile +def test_profile( + precision: Precision, scorecard_path: ScorecardProfilePath, device: ScorecardDevice +) -> None: + skip_invalid_runtime_device(Model, scorecard_path.runtime, device) + try: + profile_via_export( + profile_model, + MODEL_ID, + Model.from_pretrained(), + precision, + scorecard_path, + device, + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +@pytest.mark.parametrize( + ("precision", "scorecard_path", "device"), + get_evaluation_parameterized_pytest_config( + MODEL_ID, + EVAL_DEVICE, + ENABLED_PRECISION_RUNTIMES, + PASSING_PRECISION_RUNTIMES, + ), + ids=pytest_device_idfn, +) +@pytest.mark.inference +def test_inference( + precision: Precision, scorecard_path: ScorecardProfilePath, device: ScorecardDevice +) -> None: + skip_invalid_runtime_device(Model, scorecard_path.runtime, device) + try: + if HAS_EVAL_DATASET: + on_device_inference_for_accuracy_validation( + Model, + Model.get_eval_dataset_classes()[0], + MODEL_ID, + precision, + scorecard_path, + device, + ) + else: + inference_via_export( + inference_model, + MODEL_ID, + Model.from_pretrained(), + precision, + scorecard_path, + device, + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +@pytest.mark.inference +def test_val_data_torch() -> None: + if not HAS_EVAL_DATASET: + return + torch_inference_for_accuracy_validation( + Model.from_pretrained(), Model.get_eval_dataset_classes()[0], MODEL_ID + ) + + +@pytest.fixture(scope="module") +def torch_val_outputs() -> list[np.ndarray]: + """ + Because the below method downloads a dataset over the internet, + it is called in a fixture so it can be reused. + """ + if not HAS_EVAL_DATASET: + return [] + return torch_inference_for_accuracy_validation_outputs(MODEL_ID) + + +@pytest.fixture(scope="module") +def torch_evaluate_mock_outputs( + torch_val_outputs: list[np.ndarray], +) -> list[torch.Tensor | tuple[torch.Tensor, ...]]: + """ + Because the below method does some memory movement, + it is called in a fixture so its output can be reused. + """ + if not HAS_EVAL_DATASET: + return [] + return split_and_group_accuracy_validation_output_batches(torch_val_outputs) + + +@pytest.mark.inference +def test_torch_accuracy( + torch_evaluate_mock_outputs: list[torch.Tensor | tuple[torch.Tensor, ...]], +) -> None: + if not HAS_EVAL_DATASET: + return + torch_accuracy_on_dataset( + Model.from_pretrained(), + Model.get_eval_dataset_classes()[0], + torch_evaluate_mock_outputs, + MODEL_ID, + ) + + +@pytest.mark.parametrize( + "precision", + get_quantize_parameterized_pytest_config( + MODEL_ID, ENABLED_PRECISION_RUNTIMES, PASSING_PRECISION_RUNTIMES + ), + ids=pytest_device_idfn, +) +@pytest.mark.inference +def test_sim_accuracy( + precision: Precision, + torch_evaluate_mock_outputs: list[torch.Tensor | tuple[torch.Tensor, ...]], +) -> None: + if not HAS_EVAL_DATASET: + return + try: + sim_accuracy_on_dataset( + Model.from_pretrained(**get_model_kwargs(Model, dict(precision=precision))), + Model.get_eval_dataset_classes()[0], + MODEL_ID, + precision, + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +@pytest.mark.parametrize( + ("precision", "scorecard_path", "device"), + get_evaluation_parameterized_pytest_config( + MODEL_ID, + EVAL_DEVICE, + ENABLED_PRECISION_RUNTIMES, + PASSING_PRECISION_RUNTIMES, + ), + ids=pytest_device_idfn, +) +@pytest.mark.compute_device_accuracy +def test_val_accuracy( + precision: Precision, + scorecard_path: ScorecardProfilePath, + device: ScorecardDevice, + torch_val_outputs: list[np.ndarray], + torch_evaluate_mock_outputs: list[torch.Tensor | tuple[torch.Tensor, ...]], +) -> None: + try: + if HAS_EVAL_DATASET: + accuracy_on_dataset_via_evaluate_and_export( + export_model, + Model.from_pretrained( + **get_model_kwargs(Model, dict(precision=precision)) + ), + Model.get_eval_dataset_classes()[0], + torch_val_outputs, + torch_evaluate_mock_outputs, + MODEL_ID, + precision, + scorecard_path, + device, + ) + else: + accuracy_on_sample_inputs_via_export( + export_model, + MODEL_ID, + Model.from_pretrained(), + precision, + scorecard_path, + device, + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +@pytest.mark.parametrize( + ("precision", "scorecard_path", "device"), + get_export_parameterized_pytest_config( + MODEL_ID, + EVAL_DEVICE, + ENABLED_PRECISION_RUNTIMES, + PASSING_PRECISION_RUNTIMES, + ), + ids=pytest_device_idfn, +) +@pytest.mark.export +def test_export( + precision: Precision, scorecard_path: ScorecardProfilePath, device: ScorecardDevice +) -> None: + skip_invalid_runtime_device(Model, scorecard_path.runtime, device) + try: + export_test_e2e( + export_model, Model, MODEL_ID, precision, scorecard_path, device + ) + except CachedScorecardJobError as e: + pytest.skip(str(e)) + + +# Cache serialize() and hub.upload_model() across the module so the same +# (component, graph, input_spec) is serialized once and the resulting bytes are +# uploaded once -- matters most for multi-GB AIMET LLM bundles. +@pytest.fixture(scope="module", autouse=True) +def cached_serialize_for_export( + tmp_path_factory: pytest.TempPathFactory, +) -> Generator[pytest.MonkeyPatch, None, None]: + cache_dir = tmp_path_factory.mktemp("serialize_cache") + with pytest.MonkeyPatch.context() as mp: + model_cache: dict[str, Path] = {} + upload_cache: dict[str, hub.Model] = {} + + real_upload_model = hub.upload_model + + def _cached_upload_model( + model: hub.client.SourceModel | str, + name: str | None = None, + project: str | hub.client.Project | None = None, + ) -> hub.Model: + key = str(model) + cached = upload_cache.get(key) + if cached is None: + cached = real_upload_model(model, name, project) + upload_cache[key] = cached + return cached + + mp.setattr(hub, "upload_model", _cached_upload_model) + serialize = Model.serialize + + def _cached_serialize( + self: Model, + output_dir: str | os.PathLike, + input_spec: InputSpec | None = None, + ) -> Path: + model_key = str(input_spec) + cached = model_cache.get(model_key) + if not cached: + cached = serialize(self, cache_dir, input_spec) + model_cache[model_key] = cached + return cached + + mp.setattr(Model, "serialize", _cached_serialize) + yield mp