models.py Source File

models.py Source File#

Mobilint SDK qb Compiler: models.py Source File
Mobilint SDK qb Compiler v1.2
MCS002-EN
"""Auto-generated Pydantic models from config_schema.yaml."""
 
from __future__ import annotations
from typing import Any, Dict, List, Optional, Union
from pydantic import BaseModel, Field, ConfigDict
import yaml
import json
from pathlib import Path
 
SCHEMA_VERSION = "1.0.0"
 
 
class Uint8InputConfig(BaseModel):
    """
    @brief Configuration for uint8 input handling
 
    @details Defines whether inputs should be treated as uint8 and which specific inputs to apply this to.
 
    @param apply bool. If true, treat specified inputs as uint8
    @param inputs List[str]. List of input names to treat as uint8. If empty and apply is true, applies to all inputs
    @param division_factor float. Division factor for uint8 to float conversion (e.g., 255.0 for [0,1], 127.5 for [0,2])
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(
        default=False, description="If true, treat specified inputs as uint8"
    )
    inputs: List[str] = Field(
        default=[],
        description="List of input names to treat as uint8. If empty and apply is true, applies to all inputs",
    )
    division_factor: float = Field(
        default=255.0,
        alias="divisionFactor",
        description="Division factor for uint8 to float conversion (e.g., 255.0 for [0,1], 127.5 for [0,2])",
    )
 
    def with_updates(self, **kwargs) -> "Uint8InputConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class PreprocessingConfig(BaseModel):
    """
    @brief Configuration for input preprocessing pipeline
 
    @details Defines preprocessing operations to be applied to model inputs,
    including operations like resize, normalize, color conversion, etc.
 
    @param apply bool. If true, apply preprocessing pipeline
    @param auto_convert_format bool. If true, automatically convert input format
    @param pipeline List[Dict[str, Any]]. List of preprocessing operations to apply globally
    @param input_configs Dict[str, Any]. Per-input preprocessing configurations. Keys are input names
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(
        default=False, description="If true, apply preprocessing pipeline"
    )
    auto_convert_format: bool = Field(
        default=False,
        alias="autoConvertFormat",
        description="If true, automatically convert input format",
    )
    pipeline: Any = Field(
        default=[], description="List of preprocessing operations to apply globally"
    )
    input_configs: Any = Field(
        default={},
        alias="inputConfigs",
        description="Per-input preprocessing configurations. Keys are input names",
    )
 
    def with_updates(self, **kwargs) -> "PreprocessingConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class ResourceManagementConfig(BaseModel):
    """
    @brief Configuration for resource management during model compilation
 
    @details Controls GPU and memory management settings during the compilation process.
 
    @param weight_dtype str. Weight data type for calibration (e.g., 'float32', 'float16')
    @param use_gpu_only_for_calibration bool. If True, use GPU only during the calibration phase
    @param weight_memory WeightMemory. Weight memory management configuration
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    weight_dtype: str = Field(
        default="float32",
        alias="weightDtype",
        description="Weight data type for calibration (e.g., 'float32', 'float16')",
    )
    use_gpu_only_for_calibration: bool = Field(
        default=True,
        alias="useGPUOnlyForCalibration",
        description="If True, use GPU only during the calibration phase",
    )
 
    class WeightMemory(BaseModel):
        """
        @brief Weight memory management configuration
 
        @param method int. Weight memory management method index:<br>
            0: DeleteFloat - Delete float weights after quantization.<br>
            1: SaveFloat - Save float weights to disk.<br>
            2: MoveFloat - Move float weights to CPU.<br>
            3: KeepFloat - Keep float weights in memory.<br>
            4: KeepAll - Keep all weights in memory.<br>
        """
 
        method_list: List[str] = Field(
            default=["DeleteFloat", "SaveFloat", "MoveFloat", "KeepFloat", "KeepAll"],
            alias="methodList",
        )
        method: int = Field(default=0, alias="method")
 
    weight_memory: WeightMemory = Field(
        default_factory=WeightMemory, alias="weightMemory"
    )
 
    def with_updates(self, **kwargs) -> "ResourceManagementConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class CalibrationConfig(BaseModel):
    """
    @brief Configuration for calibration during quantization
 
    @details Defines calibration and quantization parameterization used to derive activation/weight scales
    and related statistics during quantized compilation.
 
    @param method int. Calibration method index:<br>
        0: WChALayer - Weight per-channel, Activation per-layer.<br>
        1: WChAMulti - Weight per-channel, Activation multi-layer.<br>
        2: WChALayerZeropoint - Weight per-channel, Activation per-layer with zeropoint.<br>
        3: WChAMultiZeropoint - Weight per-channel, Activation multi-layer with zeropoint.<br>
    @param output int. Output quantization type index:<br>
        0: Layer - Per-layer quantization.<br>
        1: Ch - Per-channel quantization.<br>
        2: Sigmoid - Sigmoid-based quantization.<br>
    @param mode int. Quantization mode index:<br>
        0: Max - Maximum value calibration.<br>
        1: MaxPercentile - Maximum percentile calibration.<br>
        2: Histogram - Histogram-based calibration.<br>
    @param clustering_methods int. LUT clustering method:<br>
        0: Scale - Scale-based 1D clustering.<br>
        1: Scale2D - Scale-based 2D clustering.<br>
        2: ReparameterizedL2 - Reparameterized L2 distance clustering.<br>
        3: SubgraphIoU - Subgraph intersection-over-union clustering.<br>
        4: DomainAwareL2 - Domain-aware L2 distance on common absolute grid.<br>
        5: OverlapWeightedL2 - L2 weighted by domain overlap (IoU).<br>
        6: Fast - Minimal search space for fast LUT optimization.<br>
    @param act_scale_min float. Minimum allowed activation scale (lower bound clamp)
    @param act16_scale_min float. Minimum 16-bit activation scale (actScaleMin / 256)
    @param weight_scale_min float. Minimum allowed weight scale (lower bound clamp)
    @param weight16_scale_min float. Minimum 16-bit weight scale (weightScaleMin / 256)
    @param min_clip_ratio float. Minimum clip ratio constraint applied during calibration
    @param max_calib_data_size int. Maximum number of calibration samples kept after loading or generation
    @param max_percentile MaxPercentile. MaxPercentile mode configuration
    @param fast_dist FastDist. Fast distribution calibration configuration
    @param histogram Histogram. Histogram-based calibration configuration
    @param layer_overrides LayerOverrides. Layer-specific override settings for calibration
    @param statistics Statistics. Statistics save/load configuration with percentile selection
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    method_list: List[str] = Field(
        default=["WChALayer", "WChAMulti", "WChALayerZeropoint", "WChAMultiZeropoint"],
        alias="methodList",
    )
    method: int = Field(default=1, alias="method")
    output_list: List[str] = Field(
        default=["Layer", "Ch", "Sigmoid"], alias="outputList"
    )
    output: int = Field(default=0, alias="output")
    mode_list: List[str] = Field(
        default=["Max", "MaxPercentile", "Histogram"], alias="modeList"
    )
    mode: int = Field(default=1, alias="mode")
    clustering_methods_list: List[str] = Field(
        default=[
            "Scale",
            "Scale2D",
            "ReparameterizedL2",
            "SubgraphIoU",
            "DomainAwareL2",
            "OverlapWeightedL2",
            "Fast",
        ],
        alias="clusteringMethodsList",
    )
    clustering_methods: List[int] = Field(default=[0], alias="clusteringMethods")
 
    act_scale_min: float = Field(
        default=0.0005,
        alias="actScaleMin",
        description="Minimum allowed activation scale (lower bound clamp)",
        ge=0,
        le=1,
    )
    act16_scale_min: float = Field(
        default=1.953125e-06,
        alias="act16ScaleMin",
        description="Minimum 16-bit activation scale (actScaleMin / 256)",
    )
    weight_scale_min: float = Field(
        default=1e-06,
        alias="weightScaleMin",
        description="Minimum allowed weight scale (lower bound clamp)",
        ge=0,
        le=1,
    )
    weight16_scale_min: float = Field(
        default=3.90625e-09,
        alias="weight16ScaleMin",
        description="Minimum 16-bit weight scale (weightScaleMin / 256)",
    )
    min_clip_ratio: float = Field(
        default=-1,
        alias="minClipRatio",
        description="Minimum clip ratio constraint applied during calibration",
        ge=-1,
        le=1,
    )
    max_calib_data_size: int = Field(
        default=-1,
        alias="maxCalibDataSize",
        description="Maximum number of calibration samples kept after loading or generation",
        ge=-1,
    )
 
    class MaxPercentile(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief MaxPercentile mode configuration
 
        @param percentile float. Percentile value for maxPercentile mode
        @param topk_ratio float. Top-k ratio used in maxPercentile mode
        @param max_each int. Maximum number of samples processed per iteration
        @param max_total int. Total maximum number of samples
        @param per_ch_divisor int. Divisor for per-channel buffer capacity (bufferCap = max(maxTotal / perChDivisor, maxEach))
        """
        percentile: float = Field(
            default=0.9999, description="Percentile value for maxPercentile mode"
        )
        topk_ratio: float = Field(
            default=0.01,
            alias="topKRatio",
            description="Top-k ratio used in maxPercentile mode",
        )
        max_each: int = Field(
            default=128,
            alias="maxEach",
            description="Maximum number of samples processed per iteration",
        )
        max_total: int = Field(
            default=65536,
            alias="maxTotal",
            description="Total maximum number of samples",
        )
        per_ch_divisor: int = Field(
            default=16,
            alias="perChDivisor",
            description="Divisor for per-channel buffer capacity (bufferCap = max(maxTotal / perChDivisor, maxEach))",
            ge=1,
        )
 
    max_percentile: MaxPercentile = Field(
        default_factory=MaxPercentile, alias="maxPercentile"
    )
 
    class FastDist(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Fast distribution calibration configuration
 
        @param size_cali int. 
        @param kernel_size int. 
        @param stack_size int. 
        """
        size_cali: int = Field(default=100, alias="sizeCali")
        kernel_size: int = Field(default=9, alias="kernelSize")
        stack_size: int = Field(default=32768, alias="stackSize")
 
    fast_dist: FastDist = Field(default_factory=FastDist, alias="fastDist")
 
    class Histogram(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Histogram-based calibration configuration
 
        @param search_type int. Search type for histogram calibration:<br>
            0: Percentile.<br>
            1: MSE.<br>
            2: KL.<br>
        @param percentile float. Percentile value for histogram calibration
        @param use_gpu bool. Use GPU for histogram computation
        @param num_bins int. Number of bins for histogram
        @param num_samples int. Number of samples for histogram calibration
        @param buffer_size int. Buffer size for histogram computation (-1 for auto)
        @param min_bin_width float. Minimum bin width for histogram
        @param search_percentile_min float. Minimum search percentile
        @param search_percentile_max float. Maximum search percentile
        @param num_search int. Number of search iterations
        """
        search_type_list: List[str] = Field(
            default=["Percentile", "MSE", "KL"], alias="searchTypeList"
        )
        search_type: int = Field(default=0, alias="searchType")
        percentile: float = Field(
            default=0.9999, description="Percentile value for histogram calibration"
        )
        use_gpu: bool = Field(
            default=True,
            alias="useGPU",
            description="Use GPU for histogram computation",
        )
        num_bins: int = Field(
            default=256, alias="numBins", description="Number of bins for histogram"
        )
        num_samples: int = Field(
            default=128,
            alias="numSamples",
            description="Number of samples for histogram calibration",
        )
        buffer_size: int = Field(
            default=-1,
            alias="bufferSize",
            description="Buffer size for histogram computation (-1 for auto)",
        )
        min_bin_width: float = Field(
            default=1e-06,
            alias="minBinWidth",
            description="Minimum bin width for histogram",
        )
        search_percentile_min: float = Field(
            default=0.9999,
            alias="searchPercentileMin",
            description="Minimum search percentile",
        )
        search_percentile_max: float = Field(
            default=1.0,
            alias="searchPercentileMax",
            description="Maximum search percentile",
        )
        num_search: int = Field(
            default=128, alias="numSearch", description="Number of search iterations"
        )
 
    histogram: Histogram = Field(default_factory=Histogram, alias="histogram")
 
    class LayerOverrides(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Layer-specific override settings for calibration
 
        @param act_scale_min dict. Per-layer activation scale minimum overrides. Keys are actScaleMin values (e.g. '0.0005'), values are lists of layer names to apply that override. (e.g. {'0.0005' : ['layer1', 'layer2']})
        @param percentile dict. Per-layer maxPercentile percentile overrides. Keys are layer names, values are percentile floats. (e.g. {'/model/layer0/conv': 0.999, '/model/layer5/attn': 0.99})
        """
        act_scale_min: Any = Field(
            default={},
            alias="actScaleMin",
            description="Per-layer activation scale minimum overrides. Keys are actScaleMin values (e.g. '0.0005'), values are lists of layer names to apply that override. (e.g. {'0.0005' : ['layer1', 'layer2']})",
        )
        percentile: Any = Field(
            default={},
            description="Per-layer maxPercentile percentile overrides. Keys are layer names, values are percentile floats. (e.g. {'/model/layer0/conv': 0.999, '/model/layer5/attn': 0.99})",
        )
        method: Any = Field(
            default={},
            description="Per-layer quantization method overrides. Keys are layer names, values are method ints (0=WChALayer, 1=WChAMulti, 2=WChALayerZeropoint asym, 3=WChAMultiZeropoint asym). (e.g. {'/model/layer0/conv': 3})",
        )
 
    layer_overrides: LayerOverrides = Field(
        default_factory=LayerOverrides, alias="layerOverrides"
    )
 
    class Statistics(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Statistics save/load configuration with percentile selection
 
        @param apply bool. Enable statistics save/load
        @param save_path str. Path to save statistics. If empty, not saved
        @param load_path str. Path to load statistics. If empty, not loaded
        @param percentiles List[float]. List of percentile candidates
        @param percentile_index int. Index into percentiles list to select active percentile
        """
        apply: bool = Field(default=False, description="Enable statistics save/load")
        save_path: str = Field(
            default="",
            alias="savePath",
            description="Path to save statistics. If empty, not saved",
        )
        load_path: str = Field(
            default="",
            alias="loadPath",
            description="Path to load statistics. If empty, not loaded",
        )
        percentiles: List[float] = Field(
            default=[0.9999, 0.999, 0.99, 0.9],
            description="List of percentile candidates",
        )
        percentile_index: int = Field(
            default=0,
            alias="percentileIndex",
            description="Index into percentiles list to select active percentile",
        )
 
    statistics: Statistics = Field(default_factory=Statistics, alias="statistics")
 
    def with_updates(self, **kwargs) -> "CalibrationConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class BitConfig(BaseModel):
    """
    @brief Configuration for bit precision
 
    @details Defines bit-width parameterization for activations and weights used in
    mixed-precision quantization (e.g., attention and FFN components).
 
    @param transformer Transformer. Transformer-specific bit-width configuration
    @param save_info SaveInfo. Bit allocation save/load configuration
    @param layer_overrides LayerOverrides. Layer-specific bit-width override settings
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    class Transformer(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Transformer-specific bit-width configuration
 
        @param activation Activation. Activation bit-widths for transformer components
        @param weight Weight. Weight bit-widths for transformer components
        @param mixed_precision MixedPrecision. Mixed precision configuration
        """
 
        class Activation(BaseModel):
            """
            @brief Activation bit-widths for transformer components
 
            @param query int. Query activation bit-width
            @param key int. Key activation bit-width
            @param value int. Value activation bit-width
            @param output int. Output activation bit-width
            @param ffn int. FFN activation bit-width
            @param head int. Head activation bit-width
            """
 
            query: int = Field(default=8, description="Query activation bit-width")
            key: int = Field(default=8, description="Key activation bit-width")
            value: int = Field(default=8, description="Value activation bit-width")
            output: int = Field(default=16, description="Output activation bit-width")
            ffn: int = Field(default=16, description="FFN activation bit-width")
            head: int = Field(default=8, description="Head activation bit-width")
 
        activation: Activation = Field(default_factory=Activation, alias="activation")
 
        class Weight(BaseModel):
            """
            @brief Weight bit-widths for transformer components
 
            @param query int. Query weight bit-width
            @param key int. Key weight bit-width
            @param value int. Value weight bit-width
            @param output int. Output weight bit-width
            @param ffn int. FFN weight bit-width
            @param head int. Head weight bit-width
            """
 
            query: int = Field(default=8, description="Query weight bit-width")
            key: int = Field(default=8, description="Key weight bit-width")
            value: int = Field(default=8, description="Value weight bit-width")
            output: int = Field(default=8, description="Output weight bit-width")
            ffn: int = Field(default=8, description="FFN weight bit-width")
            head: int = Field(default=8, description="Head weight bit-width")
 
        weight: Weight = Field(default_factory=Weight, alias="weight")
 
        class MixedPrecision(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief Mixed precision configuration
 
            @param weight Weight. Mixed precision configuration for weights
            @param activation Activation. Per-layer activation mixed precision configuration
            """
 
            class Weight(BaseModel):
                model_config = ConfigDict(populate_by_name=True)
 
                """
                @brief Mixed precision configuration for weights
 
                @param apply bool. If true, apply mixed-precision according to the specified bit-widths
                @param type_wise bool. Apply type-wise mixed precision
                @param prune float. Pruning ratio
                @param bit_2 float. Ratio of 2-bit quantization
                @param bit_4 float. Ratio of 4-bit quantization
                @param bit_8 float. Ratio of 8-bit quantization
                @param importance_threshold_low float. Low importance threshold
                @param importance_threshold_high float. High importance threshold
                """
                apply: bool = Field(
                    default=False,
                    description="If true, apply mixed-precision according to the specified bit-widths",
                )
                type_wise: bool = Field(
                    default=True,
                    alias="typeWise",
                    description="Apply type-wise mixed precision",
                )
                prune: float = Field(default=0, description="Pruning ratio")
                bit_2: float = Field(
                    default=0, alias="bit2", description="Ratio of 2-bit quantization"
                )
                bit_4: float = Field(
                    default=0, alias="bit4", description="Ratio of 4-bit quantization"
                )
                bit_8: float = Field(
                    default=1, alias="bit8", description="Ratio of 8-bit quantization"
                )
                importance_threshold_low: float = Field(
                    default=-1,
                    alias="importanceThreshold_low",
                    description="Low importance threshold",
                )
                importance_threshold_high: float = Field(
                    default=-1,
                    alias="importanceThreshold_high",
                    description="High importance threshold",
                )
 
            weight: Weight = Field(default_factory=Weight, alias="weight")
 
            class Activation(BaseModel):
                model_config = ConfigDict(populate_by_name=True)
 
                """
                @brief Per-layer activation mixed precision configuration
 
                @param apply bool. If true, apply per-layer activation mixed precision
                @param ratio_16bit float. Ratio of layers assigned 16-bit (used when importanceThreshold < 0)
                @param importance_threshold float. Normalized importance threshold for 16-bit assignment (negative = use ratio_16bit)
                @param search_range int. Target layer range: -1=all layers, N=first N layers
                """
                apply: bool = Field(
                    default=False,
                    description="If true, apply per-layer activation mixed precision",
                )
                ratio_16bit: float = Field(
                    default=0.45,
                    alias="ratio16Bit",
                    description="Ratio of layers assigned 16-bit (used when importanceThreshold < 0)",
                )
                importance_threshold: float = Field(
                    default=-1,
                    alias="importanceThreshold",
                    description="Normalized importance threshold for 16-bit assignment (negative = use ratio_16bit)",
                )
                search_range: int = Field(
                    default=-1,
                    alias="searchRange",
                    description="Target layer range: -1=all layers, N=first N layers",
                )
 
            activation: Activation = Field(
                default_factory=Activation, alias="activation"
            )
 
        mixed_precision: MixedPrecision = Field(
            default_factory=MixedPrecision, alias="mixedPrecision"
        )
 
    transformer: Transformer = Field(default_factory=Transformer, alias="transformer")
 
    class SaveInfo(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Bit allocation save/load configuration
 
        @param save_path str. Path to save the bit allocation. If empty, not saved
        @param load_path str. Path to load the bit allocation. If empty, not loaded
        """
        save_path: str = Field(
            default="",
            alias="savePath",
            description="Path to save the bit allocation. If empty, not saved",
        )
        load_path: str = Field(
            default="",
            alias="loadPath",
            description="Path to load the bit allocation. If empty, not loaded",
        )
 
    save_info: SaveInfo = Field(default_factory=SaveInfo, alias="saveInfo")
 
    class LayerOverrides(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Layer-specific bit-width override settings
 
        @param activation_16bits list[string]. Layer names to force 16-bit activations
        @param weight_16bits list[string]. Layer names to force 16-bit weights
        """
        activation_16bits: List[str] = Field(
            default=[],
            alias="activation16Bits",
            description="Layer names to force 16-bit activations",
        )
        weight_16bits: List[str] = Field(
            default=[],
            alias="weight16Bits",
            description="Layer names to force 16-bit weights",
        )
 
    layer_overrides: LayerOverrides = Field(
        default_factory=LayerOverrides, alias="layerOverrides"
    )
 
    def with_updates(self, **kwargs) -> "BitConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class OptqConfig(BaseModel):
    """
    @brief Configuration for OPTQ algorithm
 
    @details Defines parameters controlling whether and how OPTQ is applied during quantization,
    including layer-level inclusion/exclusion lists.
 
    @param apply bool. If true, apply OPTQ
    @param attributes Attributes. OPTQ algorithm attributes
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(default=False, description="If true, apply OPTQ")
 
    class Attributes(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief OPTQ algorithm attributes
 
        @param act_order bool. If true, use activation order
        @param block_size int. Block size used for OPTQ
        @param perc_damp float. Percentage dampening factor
        @param apply_layers List[str]. Layer names to apply OPTQ. If empty, applies to all eligible layers
        @param exclude_layers List[str]. Layer names to exclude from OPTQ
        """
        act_order: bool = Field(
            default=True, alias="actOrder", description="If true, use activation order"
        )
        block_size: int = Field(
            default=128, alias="blockSize", description="Block size used for OPTQ"
        )
        perc_damp: float = Field(
            default=0.01, alias="percDamp", description="Percentage dampening factor"
        )
        apply_layers: List[str] = Field(
            default=[],
            alias="applyLayers",
            description="Layer names to apply OPTQ. If empty, applies to all eligible layers",
        )
        exclude_layers: List[str] = Field(
            default=[],
            alias="excludeLayers",
            description="Layer names to exclude from OPTQ",
        )
 
    attributes: Attributes = Field(default_factory=Attributes, alias="attributes")
 
    def with_updates(self, **kwargs) -> "OptqConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class ModConfig(BaseModel):
    """
    @brief Configuration for Minimum Output Difference algorithm
 
    @details Defines parameters controlling whether and how MOD is applied during quantization,
    including layer-level inclusion/exclusion lists.
 
    @param apply bool. If true, apply MOD
    @param attributes Attributes. MOD algorithm attributes
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(default=False, description="If true, apply MOD")
 
    class Attributes(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief MOD algorithm attributes
 
        @param epochs int. Number of training epochs
        @param warmup_epochs int. Number of warmup epochs
        @param lr_min_ratio float. Minimum learning rate ratio
        @param save_dir str. Directory to save MOD results
        @param seed int. Random seed for MOD
        @param apply_layers List[str]. Layer names to apply MOD. If empty, applies to all eligible layers
        @param exclude_layers List[str]. Layer names to exclude from MOD
        @param mod_after_layer_name str. Apply MOD after this layer
        @param anchors List. Anchor configurations for detection models. Nested list of anchor box sizes for detection models. Structure: List[List[List[int]]] where outer list is per detection head (e.g. small/medium/large), middle list is anchors per head, inner list is [width, height]. Example: [[[12,16],[19,36],[40,28]], [[36,75],[76,55],[72,146]], [[142,110],[192,243],[459,401]]]
        @param use_xyxy bool. Use XYXY format for bounding boxes
        @param learning_rates LearningRates. Learning rate configuration for MOD
        @param training Training. MOD training configuration
        @param loss Loss. MOD loss configuration
        @param post_processing PostProcessing. Post-processing configuration for detection models
        """
        epochs: int = Field(default=4, description="Number of training epochs")
        warmup_epochs: int = Field(
            default=1, alias="warmupEpochs", description="Number of warmup epochs"
        )
        lr_min_ratio: float = Field(
            default=0.0001,
            alias="lrMinRatio",
            description="Minimum learning rate ratio",
        )
        save_dir: str = Field(
            default="", alias="saveDir", description="Directory to save MOD results"
        )
        seed: int = Field(default=0, description="Random seed for MOD")
        apply_layers: List[str] = Field(
            default=[],
            alias="applyLayers",
            description="Layer names to apply MOD. If empty, applies to all eligible layers",
        )
        exclude_layers: List[str] = Field(
            default=[],
            alias="excludeLayers",
            description="Layer names to exclude from MOD",
        )
        mod_after_layer_name: str = Field(
            default="",
            alias="modAfterLayerName",
            description="Apply MOD after this layer",
        )
        anchors: Any = Field(
            default=[],
            description="Anchor configurations for detection models. Nested list of anchor box sizes for detection models. Structure: List[List[List[int]]] where outer list is per detection head (e.g. small/medium/large), middle list is anchors per head, inner list is [width, height]. Example: [[[12,16],[19,36],[40,28]], [[36,75],[76,55],[72,146]], [[142,110],[192,243],[459,401]]]",
        )
        use_xyxy: bool = Field(
            default=False,
            alias="useXYXY",
            description="Use XYXY format for bounding boxes",
        )
 
        class LearningRates(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief Learning rate configuration for MOD
 
            @param act_scale float. Learning rate for activation scale
            @param zeropoint float. Learning rate for zeropoint
            @param weight_scale float. Learning rate for weight scale
            @param weight float. Learning rate for weight
            @param bias float. Learning rate for bias
            """
            act_scale: float = Field(
                default=0.0,
                alias="actScale",
                description="Learning rate for activation scale",
            )
            zeropoint: float = Field(
                default=0.0, description="Learning rate for zeropoint"
            )
            weight_scale: float = Field(
                default=0.0,
                alias="weightScale",
                description="Learning rate for weight scale",
            )
            weight: float = Field(default=4e-06, description="Learning rate for weight")
            bias: float = Field(default=4e-06, description="Learning rate for bias")
 
        learning_rates: LearningRates = Field(
            default_factory=LearningRates, alias="learningRates"
        )
 
        class Training(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief MOD training configuration
 
            @param batch_size int. Batch size for MOD training
            @param q_drop float. Quantization drop probability
            @param quantize_weight bool. Whether to quantize weights
            @param weight_scale_init str. Weight scale initialization method
            @param downresol_mode str. Downresolution mode
            @param scheduler_type str. LR scheduler type
            """
            batch_size: int = Field(
                default=1, alias="batchSize", description="Batch size for MOD training"
            )
            q_drop: float = Field(
                default=0.0, alias="qDrop", description="Quantization drop probability"
            )
            quantize_weight: bool = Field(
                default=True,
                alias="quantizeWeight",
                description="Whether to quantize weights",
            )
            weight_scale_init: str = Field(
                default="MinMax",
                alias="weightScaleInit",
                description="Weight scale initialization method",
            )
            downresol_mode: str = Field(
                default="STE", alias="downresolMode", description="Downresolution mode"
            )
            scheduler_type: str = Field(
                default="Cosine", alias="schedulerType", description="LR scheduler type"
            )
 
        training: Training = Field(default_factory=Training, alias="training")
 
        class Loss(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief MOD loss configuration
 
            @param type str. Loss type (MSE, KL, etc.)
            @param use_outputs bool. Use model outputs for loss computation
            @param kl_temperature float. KL divergence temperature
            @param recon_prob float. Reconstruction probability
            @param recon_coeff float. Reconstruction coefficient
            @param lambda_0 float. Loss weight lambda_0
            @param lambda_1 float. Loss weight lambda_1
            @param lambda_2 float. Loss weight lambda_2
            @param lambda_3 float. Loss weight lambda_3
            @param custom_loss_jit_path str. Path to custom JIT-compiled loss function. Refer to /workspace/quantizer/pyutils/mel.pt
            """
            type: str = Field(default="MSE", description="Loss type (MSE, KL, etc.)")
            use_outputs: bool = Field(
                default=False,
                alias="useOutputs",
                description="Use model outputs for loss computation",
            )
            kl_temperature: float = Field(
                default=1.0,
                alias="KLTemperature",
                description="KL divergence temperature",
            )
            recon_prob: float = Field(
                default=1.0, alias="reconProb", description="Reconstruction probability"
            )
            recon_coeff: float = Field(
                default=1.0,
                alias="reconCoeff",
                description="Reconstruction coefficient",
            )
            lambda_0: float = Field(
                default=1.0, alias="lambda0", description="Loss weight lambda_0"
            )
            lambda_1: float = Field(
                default=1.0, alias="lambda1", description="Loss weight lambda_1"
            )
            lambda_2: float = Field(
                default=1.0, alias="lambda2", description="Loss weight lambda_2"
            )
            lambda_3: float = Field(
                default=1.0, alias="lambda3", description="Loss weight lambda_3"
            )
            custom_loss_jit_path: str = Field(
                default="",
                alias="customLossJITPath",
                description="Path to custom JIT-compiled loss function. Refer to /workspace/quantizer/pyutils/mel.pt",
            )
 
        loss: Loss = Field(default_factory=Loss, alias="loss")
 
        class PostProcessing(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief Post-processing configuration for detection models
 
            @param post str. Post-processing type
            @param box_conf_thres float. Box confidence threshold
            @param box_iou_thres float. Box IoU threshold
            """
            post: str = Field(default="", description="Post-processing type")
            box_conf_thres: float = Field(
                default=0, alias="boxConfThres", description="Box confidence threshold"
            )
            box_iou_thres: float = Field(
                default=0, alias="boxIoUThres", description="Box IoU threshold"
            )
 
        post_processing: PostProcessing = Field(
            default_factory=PostProcessing, alias="postProcessing"
        )
 
    attributes: Attributes = Field(default_factory=Attributes, alias="attributes")
 
    def with_updates(self, **kwargs) -> "ModConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class LlmConfig(BaseModel):
    """
    @brief Configuration for Large Language Model (LLM) compilation
 
    @details Defines LLM-specific settings including sequence lengths, cache configurations,
    and runtime parameters for efficient LLM inference.
 
    @param apply bool. If True, apply LLM-specific configurations
    @param npu_parallel_degree int. Number of NPU partitions for FFN tensor parallelism (1 = disabled). Applied before OptimizeFFN if both active.
    @param attributes Attributes. LLM attributes configuration
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(
        default=False, description="If True, apply LLM-specific configurations"
    )
    npu_parallel_degree: int = Field(
        default=1,
        alias="npuParallelDegree",
        description="Number of NPU partitions for FFN tensor parallelism (1 = disabled). Applied before OptimizeFFN if both active.",
    )
 
    class Attributes(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief LLM attributes configuration
 
        @param max_data_length int. Maximum data length
        @param max_sequence_length int. Maximum sequence length
        @param max_cache_length int. Maximum cache length
        @param max_core_data_length int. Maximum core data length
        @param calibration Calibration. LLM calibration settings
        @param runtime Runtime. LLM runtime settings
        @param debug Debug. LLM debug settings
        """
        max_data_length: int = Field(
            default=4096, alias="maxDataLength", description="Maximum data length"
        )
        max_sequence_length: int = Field(
            default=4096,
            alias="maxSequenceLength",
            description="Maximum sequence length",
        )
        max_cache_length: int = Field(
            default=4096, alias="maxCacheLength", description="Maximum cache length"
        )
        max_core_data_length: int = Field(
            default=128,
            alias="maxCoreDataLength",
            description="Maximum core data length",
        )
 
        class Calibration(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief LLM calibration settings
 
            @param random_seq_length int. Random sequence length used for calibration
            @param use_full_seq_length bool. If True, use the full sequence length for calibration
            """
            random_seq_length: int = Field(
                default=80,
                alias="randomSeqLength",
                description="Random sequence length used for calibration",
            )
            use_full_seq_length: bool = Field(
                default=False,
                alias="useFullSeqLength",
                description="If True, use the full sequence length for calibration",
            )
 
        calibration: Calibration = Field(
            default_factory=Calibration, alias="calibration"
        )
 
        class Runtime(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief LLM runtime settings
 
            @param use_global_core bool. If True, use a global core
            @param batch_size int. Batch size
            @param npu_core_ids List[int]. List of NPU core IDs
            @param dynamic_rope bool. If True, enable dynamic RoPE (rotary position embedding)
            @param dynamic_mask bool. If True, enable dynamic mask (attention mask as runtime input)
            """
            use_global_core: bool = Field(
                default=False,
                alias="useGlobalCore",
                description="If True, use a global core",
            )
            batch_size: int = Field(
                default=1, alias="batchSize", description="Batch size"
            )
            npu_core_ids: List[int] = Field(
                default=[0], alias="npuCoreIds", description="List of NPU core IDs"
            )
            dynamic_rope: bool = Field(
                default=False,
                alias="dynamicRope",
                description="If True, enable dynamic RoPE (rotary position embedding)",
            )
            dynamic_mask: bool = Field(
                default=False,
                alias="dynamicMask",
                description="If True, enable dynamic mask (attention mask as runtime input)",
            )
 
        runtime: Runtime = Field(default_factory=Runtime, alias="runtime")
 
        class Debug(BaseModel):
            model_config = ConfigDict(populate_by_name=True)
 
            """
            @brief LLM debug settings
 
            @param apply bool. Enable LLM debug mode
            @param batch_debug_bundle_size int. Batch debug bundle size
            """
            apply: bool = Field(default=False, description="Enable LLM debug mode")
            batch_debug_bundle_size: int = Field(
                default=0,
                alias="batchDebugBundleSize",
                description="Batch debug bundle size",
            )
 
        debug: Debug = Field(default_factory=Debug, alias="debug")
 
    attributes: Attributes = Field(default_factory=Attributes, alias="attributes")
 
    def with_updates(self, **kwargs) -> "LlmConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class MoeConfig(BaseModel):
    """
    @brief Sparse MoE expert-selection configuration (calibration only)
 
    @details Controls which experts are calibrated inside SparseMoe modules. selectionMode
    is a calibration-only knob: it picks which experts collect statistics (and, for
    TopK, on which tokens). It does NOT change inference routing — the forward path
    always routes the router's top-K experts regardless of this setting.
    scoreThreshold is only used when selectionMode is Threshold.
 
    @param selection_mode int. Expert selection mode index (calibration only; inference is always top-K):<br>
        0: TopK - Calibrate only the router's top-K experts per token (matches inference routing).<br>
        1: All - Calibrate every expert on the full sequence.<br>
        2: Threshold - Calibrate all experts whose routing score exceeds scoreThreshold.<br>
    @param score_threshold float. Routing score threshold used when selectionMode is Threshold (calibration only)
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    selection_mode_list: List[str] = Field(
        default=["TopK", "All", "Threshold"], alias="selectionModeList"
    )
    selection_mode: int = Field(default=0, alias="selectionMode")
 
    score_threshold: float = Field(
        default=0.0,
        alias="scoreThreshold",
        description="Routing score threshold used when selectionMode is Threshold (calibration only)",
        ge=0,
    )
 
    def with_updates(self, **kwargs) -> "MoeConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class EquivalentTransformationConfig(BaseModel):
    """
    @brief Configuration for equivalent transformation techniques
 
    @details Defines parameters for various equivalent transformation methods including
    NormConv, QK smoothing, and rotation matrices for improved quantization.
 
    @param seed int. Random seed for transformation
    @param apply_hadamard_rotation_matrix bool. Apply Hadamard rotation matrix
    @param norm_conv NormConv. NormConv equivalent transformation
    @param qk Qk. QK smoothing transformation
    @param ud Ud. UD transformation
    @param vo Vo. VO transformation
    @param feed_forward_multi_lut FeedForwardMultiLut. Feed-forward multi-LUT transformation
    @param spin_r1 SpinR1. SpinR1 rotation transformation
    @param head_out_ch_rotation HeadOutChRotation. Head output channel rotation transformation
    @param in_rotation InRotation. Input rotation transformation
    @param spin_r2 SpinR2. SpinR2 rotation transformation
    @param qk_rotation QkRotation. QK rotation transformation
    @param flatten_quant FlattenQuant. Flatten quantization transformation
    @param optimize_ffn OptimizeFfn. FFN optimization
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    seed: int = Field(default=0, description="Random seed for transformation")
    apply_hadamard_rotation_matrix: bool = Field(
        default=True,
        alias="applyHadamardRotationMatrix",
        description="Apply Hadamard rotation matrix",
    )
 
    class NormConv(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief NormConv equivalent transformation
 
        @param apply bool. Apply NormConv transformation
        @param learn bool. Learn transformation parameters
        @param smoothing_factor float. Smoothing factor
        @param min_gamma float. Minimum gamma value
        @param max_gamma float. Maximum gamma value
        """
        apply: bool = Field(default=False, description="Apply NormConv transformation")
        learn: bool = Field(
            default=False, description="Learn transformation parameters"
        )
        smoothing_factor: float = Field(
            default=0.5, alias="smoothingFactor", description="Smoothing factor"
        )
        min_gamma: float = Field(
            default=0.0001, alias="minGamma", description="Minimum gamma value"
        )
        max_gamma: float = Field(
            default=10000.0, alias="maxGamma", description="Maximum gamma value"
        )
 
    norm_conv: NormConv = Field(default_factory=NormConv, alias="NormConv")
 
    class Qk(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief QK smoothing transformation
 
        @param apply bool. Apply QK transformation
        @param smoothing_factor float. Smoothing factor
        @param min_gamma float. Minimum gamma value
        @param max_gamma float. Maximum gamma value
        """
        apply: bool = Field(default=False, description="Apply QK transformation")
        smoothing_factor: float = Field(
            default=0.5, alias="smoothingFactor", description="Smoothing factor"
        )
        min_gamma: float = Field(
            default=0.0001, alias="minGamma", description="Minimum gamma value"
        )
        max_gamma: float = Field(
            default=10000.0, alias="maxGamma", description="Maximum gamma value"
        )
 
    qk: Qk = Field(default_factory=Qk, alias="QK")
 
    class Ud(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief UD transformation
 
        @param apply bool. Apply UD transformation
        @param learn bool. Learn transformation parameters
        @param smoothing_factor float. Smoothing factor
        @param min_gamma float. Minimum gamma value
        @param max_gamma float. Maximum gamma value
        """
        apply: bool = Field(default=False, description="Apply UD transformation")
        learn: bool = Field(
            default=False, description="Learn transformation parameters"
        )
        smoothing_factor: float = Field(
            default=0.5, alias="smoothingFactor", description="Smoothing factor"
        )
        min_gamma: float = Field(
            default=0.0001, alias="minGamma", description="Minimum gamma value"
        )
        max_gamma: float = Field(
            default=10000.0, alias="maxGamma", description="Maximum gamma value"
        )
 
    ud: Ud = Field(default_factory=Ud, alias="UD")
 
    class Vo(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief VO transformation
 
        @param apply bool. Apply VO transformation
        @param smoothing_factor float. Smoothing factor
        @param min_gamma float. Minimum gamma value
        @param max_gamma float. Maximum gamma value
        """
        apply: bool = Field(default=False, description="Apply VO transformation")
        smoothing_factor: float = Field(
            default=0.5, alias="smoothingFactor", description="Smoothing factor"
        )
        min_gamma: float = Field(
            default=0.0001, alias="minGamma", description="Minimum gamma value"
        )
        max_gamma: float = Field(
            default=10000.0, alias="maxGamma", description="Maximum gamma value"
        )
 
    vo: Vo = Field(default_factory=Vo, alias="VO")
 
    class FeedForwardMultiLut(BaseModel):
        """
        @brief Feed-forward multi-LUT transformation
 
        @param apply bool. Apply feed-forward multi-LUT transformation
        @param breakpoints List[float]. Breakpoints for multi-LUT
        """
 
        apply: bool = Field(
            default=False, description="Apply feed-forward multi-LUT transformation"
        )
        breakpoints: List[float] = Field(
            default=[-8.0, -4.0, 0], description="Breakpoints for multi-LUT"
        )
 
    feed_forward_multi_lut: FeedForwardMultiLut = Field(
        default_factory=FeedForwardMultiLut, alias="FeedForwardMultiLUT"
    )
 
    class SpinR1(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief SpinR1 rotation transformation
 
        @param apply bool. Apply SpinR1 transformation
        @param matrix_path str. Path to rotation matrix file
        """
        apply: bool = Field(default=False, description="Apply SpinR1 transformation")
        matrix_path: str = Field(
            default="", alias="matrixPath", description="Path to rotation matrix file"
        )
 
    spin_r1: SpinR1 = Field(default_factory=SpinR1, alias="SpinR1")
 
    class HeadOutChRotation(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Head output channel rotation transformation
 
        @param apply bool. Apply head output channel rotation
        @param matrix_path str. Path to rotation matrix file
        """
        apply: bool = Field(
            default=False, description="Apply head output channel rotation"
        )
        matrix_path: str = Field(
            default="", alias="matrixPath", description="Path to rotation matrix file"
        )
 
    head_out_ch_rotation: HeadOutChRotation = Field(
        default_factory=HeadOutChRotation, alias="HeadOutChRotation"
    )
 
    class InRotation(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Input rotation transformation
 
        @param apply bool. Apply input rotation
        @param matrix_path str. Path to rotation matrix file
        @param input_names List[str]. Names of the input layers to rotate
        """
        apply: bool = Field(default=False, description="Apply input rotation")
        matrix_path: str = Field(
            default="", alias="matrixPath", description="Path to rotation matrix file"
        )
        input_names: List[str] = Field(
            default=[],
            alias="inputNames",
            description="Names of the input layers to rotate",
        )
 
    in_rotation: InRotation = Field(default_factory=InRotation, alias="InRotation")
 
    class SpinR2(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief SpinR2 rotation transformation
 
        @param apply bool. Apply SpinR2 transformation
        @param learn bool. Learn rotation matrix
        @param matrix_path str. Path to rotation matrix file
        """
        apply: bool = Field(default=False, description="Apply SpinR2 transformation")
        learn: bool = Field(default=False, description="Learn rotation matrix")
        matrix_path: str = Field(
            default="", alias="matrixPath", description="Path to rotation matrix file"
        )
 
    spin_r2: SpinR2 = Field(default_factory=SpinR2, alias="SpinR2")
 
    class QkRotation(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief QK rotation transformation
 
        @param apply bool. Apply QK rotation transformation
        @param matrix_path str. Path to rotation matrix file
        """
        apply: bool = Field(
            default=False, description="Apply QK rotation transformation"
        )
        matrix_path: str = Field(
            default="", alias="matrixPath", description="Path to rotation matrix file"
        )
 
    qk_rotation: QkRotation = Field(default_factory=QkRotation, alias="QKRotation")
 
    class FlattenQuant(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief Flatten quantization transformation
 
        @param apply bool. Apply flatten quantization
        @param learn bool. Learn flattening parameters
        @param apply_threshold float. Threshold for applying flatten quantization
        @param max_overhead float. Maximum overhead allowed for flattening
        """
        apply: bool = Field(default=False, description="Apply flatten quantization")
        learn: bool = Field(default=False, description="Learn flattening parameters")
        apply_threshold: float = Field(
            default=0.33,
            alias="applyThreshold",
            description="Threshold for applying flatten quantization",
        )
        max_overhead: float = Field(
            default=0.02,
            alias="maxOverhead",
            description="Maximum overhead allowed for flattening",
        )
 
    flatten_quant: FlattenQuant = Field(
        default_factory=FlattenQuant, alias="FlattenQuant"
    )
 
    class OptimizeFfn(BaseModel):
        model_config = ConfigDict(populate_by_name=True)
 
        """
        @brief FFN optimization
 
        @param apply bool. Apply FFN optimization
        @param ch_per_ffn int. Optimize FFN split (-1 for auto)
        """
        apply: bool = Field(default=False, description="Apply FFN optimization")
        ch_per_ffn: int = Field(
            default=-1, alias="chPerFFN", description="Optimize FFN split (-1 for auto)"
        )
 
    optimize_ffn: OptimizeFfn = Field(default_factory=OptimizeFfn, alias="OptimizeFFN")
 
    def with_updates(self, **kwargs) -> "EquivalentTransformationConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class SearchWeightScaleConfig(BaseModel):
    """
    @brief Configuration for weight scale search
 
    @details Defines which transformer components should have their weight scales
    searched for optimal quantization.
 
    @param apply bool. If true, apply weight scale search
    @param transformer Transformer. Transformer components for weight scale search
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(default=False, description="If true, apply weight scale search")
 
    class Transformer(BaseModel):
        """
        @brief Transformer components for weight scale search
 
        @param query bool. Search weight scale for query
        @param key bool. Search weight scale for key
        @param value bool. Search weight scale for value
        @param out bool. Search weight scale for output
        @param ffn bool. Search weight scale for FFN
        """
 
        query: bool = Field(default=False, description="Search weight scale for query")
        key: bool = Field(default=False, description="Search weight scale for key")
        value: bool = Field(default=False, description="Search weight scale for value")
        out: bool = Field(default=False, description="Search weight scale for output")
        ffn: bool = Field(default=False, description="Search weight scale for FFN")
 
    transformer: Transformer = Field(default_factory=Transformer, alias="transformer")
 
    def with_updates(self, **kwargs) -> "SearchWeightScaleConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class LoadScaleConfig(BaseModel):
    """
        @brief QAT activation scale loading configuration
 
        @details Loads pre-trained activation scales from safetensors files and applies them
        to specified layers before scale/zeropoint computation.
        NOTE: entries is stored as raw JSON (type: list) because the generator does
        not support list[CustomStruct]. The schema records the field shape; parsing
        is done manually in applyQATLoadScaleQuantType / applyQATLoadScales.
 
        @param apply bool. If true, load and apply QAT scales from safetensors files
        @param entries List. List of file entries. Each entry is a dict:
      { path: str, scales: [ { key: str, layers: [str] } ] }
    path: safetensors file path; key: tensor name in the file;
    layers: layer names whose activation scale will be overridden.
 
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(
        default=False,
        description="If true, load and apply QAT scales from safetensors files",
    )
    entries: Any = Field(
        default=[],
        description="List of file entries. Each entry is a dict: { path: str, scales: [ { key: str, layers: [str] } ] } path: safetensors file path; key: tensor name in the file; layers: layer names whose activation scale will be overridden.",
    )
 
    def with_updates(self, **kwargs) -> "LoadScaleConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class RuntimeOptions(BaseModel):
    """
    @brief Runtime options for compilation
 
    @details Contains runtime-specific settings like version info and cache options.
 
    @param version str. Compiler version string (e.g., 0.0.0)
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    version: str = Field(
        default="0.0.0", description="Compiler version string (e.g., 0.0.0)"
    )
 
    def with_updates(self, **kwargs) -> "RuntimeOptions":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class SaveSampleConfig(BaseModel):
    """
    @brief Sample data generation and saving configuration
 
    @param apply bool. Enable sample data saving
    @param mode str. Inference mode: infer (standard) or inferWithCache (LLM cache models)
    @param batch_size int. Number of inference batches to generate
    @param batch_seq_lens List. Per-batch step-wise sequence lengths for inferWithCache mode. e.g. [[80, 1], [240, 10]]
    @param save_folder str. Output folder for sample data
    @param dtype str. Data type for saved samples: float or int8
    """
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    apply: bool = Field(default=False, description="Enable sample data saving")
    mode: str = Field(
        default="infer",
        description="Inference mode: infer (standard) or inferWithCache (LLM cache models)",
    )
    batch_size: int = Field(
        default=1,
        alias="batchSize",
        description="Number of inference batches to generate",
    )
    batch_seq_lens: Any = Field(
        default=[],
        alias="batchSeqLens",
        description="Per-batch step-wise sequence lengths for inferWithCache mode. e.g. [[80, 1], [240, 10]]",
    )
    save_folder: str = Field(
        default="sampleInout",
        alias="saveFolder",
        description="Output folder for sample data",
    )
    dtype: str = Field(
        default="float", description="Data type for saved samples: float or int8"
    )
 
    def with_updates(self, **kwargs) -> "SaveSampleConfig":
        """Return a copy with updated fields."""
        return self.model_copy(update=kwargs)
 
 
class CompileConfig(BaseModel):
    """Unified compilation configuration for Mobilint MXQ compilation."""
 
    model_config = ConfigDict(
        populate_by_name=True,
        extra="forbid",
    )
 
    model_paths: List[str] = Field(
        default=[], alias="modelPaths", description="Paths to model files"
    )
    calib_data_path: List[str] = Field(
        default=[], alias="calibDataPaths", description="Paths to calibration datasets"
    )
    save_paths: List[str] = Field(
        default=["./tmp.mxq"],
        alias="savePaths",
        description="Output MXQ filename/paths",
    )
    use_random_calib: bool = Field(
        default=False, alias="useRandomCalib", description="Use random calibration"
    )
    save_msgpack_name: Optional[str] = Field(
        default=None,
        alias="saveMsgpackName",
        description="Name of the msgpack file to save",
    )
    inference_scheme: str = Field(
        default="single", alias="inferenceScheme", description="NPU inference scheme"
    )
    cpu_offload: bool = Field(
        default=False,
        alias="cpuOffload",
        description="Enable CPU offload for unsupported operators",
    )
    force_npu_input_reposition: bool = Field(
        default=False,
        alias="forceNpuInputReposition",
        description="Force input reposition operations to run on NPU instead of CPU",
    )
    force_npu_output_reposition: bool = Field(
        default=False,
        alias="forceNpuOutputReposition",
        description="Force output reposition operations to run on NPU instead of CPU",
    )
    optimize_option: int = Field(
        default=1,
        alias="optimizeOption",
        description="Compiler optimization selector",
        ge=0,
    )
    buffer_mode: int = Field(
        default=1, alias="bufferMode", description="Buffer serialization mode"
    )
    input_shape_dict: Any = Field(
        default={}, alias="inputShapeDict", description="Dictionary of input shapes"
    )
    device: str = Field(default="gpu", description="Device for computation")
    dtype: str = Field(default="float", description="Data type for computation")
    debug: bool = Field(default=False, description="Enable debug mode")
    trace: bool = Field(default=False, description="Enable trace mode")
    image_channels: int = Field(
        default=0,
        alias="imageChannels",
        description="Number of image channels (0 for auto-detect)",
    )
    config_version: str = Field(
        default="1.0.0", alias="configVersion", description="Config schema version"
    )
    split_blocks: List[int] = Field(
        default=[],
        alias="splitBlocks",
        description="Multi-MXQ split points by transformer block index",
    )
    split_parts: int = Field(
        default=0,
        alias="splitParts",
        description="Evenly split transformer blocks into N MXQ parts",
    )
 
    uint8_input: Uint8InputConfig = Field(
        default_factory=Uint8InputConfig, alias="uint8Input"
    )
    preprocessing: PreprocessingConfig = Field(default_factory=PreprocessingConfig)
    resource_management: ResourceManagementConfig = Field(
        default_factory=ResourceManagementConfig, alias="resourceManagement"
    )
    calibration: CalibrationConfig = Field(default_factory=CalibrationConfig)
    bit: BitConfig = Field(default_factory=BitConfig)
    optq: OptqConfig = Field(default_factory=OptqConfig)
    mod: ModConfig = Field(default_factory=ModConfig)
    llm: LlmConfig = Field(default_factory=LlmConfig)
    moe: MoeConfig = Field(default_factory=MoeConfig)
    equivalent_transformation: EquivalentTransformationConfig = Field(
        default_factory=EquivalentTransformationConfig, alias="equivalentTransformation"
    )
    search_weight_scale: SearchWeightScaleConfig = Field(
        default_factory=SearchWeightScaleConfig, alias="searchWeightScale"
    )
    load_scale: LoadScaleConfig = Field(
        default_factory=LoadScaleConfig, alias="loadScale"
    )
    runtime_options: RuntimeOptions = Field(
        default_factory=RuntimeOptions, alias="runtimeOptions"
    )
    save_sample: SaveSampleConfig = Field(
        default_factory=SaveSampleConfig, alias="saveSample"
    )
 
    def with_uint8_input(self, **kwargs) -> "CompileConfig":
        """Return a copy with uint8_input settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.uint8_input.model_copy(update=data)
        return self.model_copy(update={"uint8_input": new_cfg})
 
    def with_preprocessing(self, **kwargs) -> "CompileConfig":
        """Return a copy with preprocessing settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.preprocessing.model_copy(update=data)
        return self.model_copy(update={"preprocessing": new_cfg})
 
    def with_optq(self, **kwargs) -> "CompileConfig":
        """Return a copy with optq settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.optq.model_copy(update=data)
        return self.model_copy(update={"optq": new_cfg})
 
    def with_mod(self, **kwargs) -> "CompileConfig":
        """Return a copy with mod settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.mod.model_copy(update=data)
        return self.model_copy(update={"mod": new_cfg})
 
    def with_llm(self, **kwargs) -> "CompileConfig":
        """Return a copy with llm settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.llm.model_copy(update=data)
        return self.model_copy(update={"llm": new_cfg})
 
    def with_search_weight_scale(self, **kwargs) -> "CompileConfig":
        """Return a copy with search_weight_scale settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.search_weight_scale.model_copy(update=data)
        return self.model_copy(update={"search_weight_scale": new_cfg})
 
    def with_load_scale(self, **kwargs) -> "CompileConfig":
        """Return a copy with load_scale settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.load_scale.model_copy(update=data)
        return self.model_copy(update={"load_scale": new_cfg})
 
    def with_save_sample(self, **kwargs) -> "CompileConfig":
        """Return a copy with save_sample settings enabled."""
        data = {"apply": True, **kwargs}
        new_cfg = self.save_sample.model_copy(update=data)
        return self.model_copy(update={"save_sample": new_cfg})
 
    @classmethod
    def from_file(cls, path: Union[str, Path]) -> "CompileConfig":
        """Load config from YAML or JSON file."""
        path = Path(path)
        with open(path) as f:
            if path.suffix in (".yaml", ".yml"):
                data = yaml.safe_load(f)
            else:
                data = json.load(f)
        data = cls._flatten_grouped_json(data)
        return cls.model_validate(data)
 
    @staticmethod
    def _flatten_grouped_json(data: dict) -> dict:
        """Flatten grouped JSON keys (e.g. quantization.calibration) to flat structure."""
        data = data.copy()
        if "quantization" in data:
            group = data.pop("quantization")
            if "calibration" in group:
                data["calibration"] = group["calibration"]
            if "bit" in group:
                data["bit"] = group["bit"]
        if "advancedQuantization" in data:
            group = data.pop("advancedQuantization")
            if "optq" in group:
                data["optq"] = group["optq"]
            if "mod" in group:
                data["mod"] = group["mod"]
            if "EquivalentTransformation" in group:
                data["equivalentTransformation"] = group["EquivalentTransformation"]
            if "searchWeightScale" in group:
                data["searchWeightScale"] = group["searchWeightScale"]
            if "loadScale" in group:
                data["loadScale"] = group["loadScale"]
        return data
 
    @classmethod
    def from_preset(cls, name: str) -> "CompileConfig":
        """Load config from a preset."""
        from .presets import get_preset
 
        return get_preset(name)
 
    @staticmethod
    def _group_to_json(data: dict) -> dict:
        """Group flat keys back into nested JSON structure (inverse of _flatten_grouped_json)."""
        data = data.copy()
        quantization_group = {}
        if "calibration" in data:
            quantization_group["calibration"] = data.pop("calibration")
        if "bit" in data:
            quantization_group["bit"] = data.pop("bit")
        if quantization_group:
            data["quantization"] = quantization_group
        advancedQuantization_group = {}
        if "optq" in data:
            advancedQuantization_group["optq"] = data.pop("optq")
        if "mod" in data:
            advancedQuantization_group["mod"] = data.pop("mod")
        if "equivalentTransformation" in data:
            advancedQuantization_group["EquivalentTransformation"] = data.pop(
                "equivalentTransformation"
            )
        if "searchWeightScale" in data:
            advancedQuantization_group["searchWeightScale"] = data.pop(
                "searchWeightScale"
            )
        if "loadScale" in data:
            advancedQuantization_group["loadScale"] = data.pop("loadScale")
        if advancedQuantization_group:
            data["advancedQuantization"] = advancedQuantization_group
        return data
 
    def to_file(self, path: Union[str, Path]) -> None:
        """Save config to YAML or JSON file."""
        path = Path(path)
        data = self.model_dump(by_alias=True, exclude_none=True)
        data = self._group_to_json(data)
        with open(path, "w") as f:
            if path.suffix in (".yaml", ".yml"):
                yaml.dump(data, f, default_flow_style=False)
            else:
                json.dump(data, f, indent=2)