llmcompressor.entrypoints.utils

def post_process(
    model_args: ModelArguments | None = None,
    recipe_args: RecipeArguments | None = None,
    output_dir: str | None = None,
):
    """
    Saves the model and tokenizer/processor to the output directory if model_args,
    output_dir is provided.

    If the `output_dir` is not the default directory, the method resets lifecycle
    actions. The model is saved in a compressed format if specified in `model_args`.
    Additionally, the tokenizer or processor, if available, is also saved.

    Raises:
        ValueError: If saving fails due to an invalid `output_dir` or other issues.
    """
    if model_args is not None and output_dir is not None:
        if recipe_args is not None and getattr(recipe_args, "stage", None) is not None:
            output_dir = os.path.join(output_dir, recipe_args.stage)
            os.makedirs(output_dir, exist_ok=True)
            logger.info(f"[Save] Stage detected. Updating output_dir to {output_dir}")

        # TODO: support general saving parameters, beyond save_compressed
        model_args.model.save_pretrained(
            output_dir, save_compressed=model_args.save_compressed
        )

        if model_args.processor is not None:
            model_args.processor.save_pretrained(output_dir)

    else:
        logger.warning(
            "Optimized model is not saved. To save, please provide"
            "`output_dir` as input arg."
            "Ex. `oneshot(..., output_dir=...)`"
        )

    # Reset the one-time-use session upon completion
    if recipe_args is not None and recipe_args.clear_sparse_session:
        reset_session()

def pre_process(
    model_args: ModelArguments,
    dataset_args: DatasetArguments,
    output_dir: str | None,
):
    """
    Prepares the model and tokenizer/processor for calibration.
    - Initializes the model if it's specified as a path or string.
    - Applies patches to fix tied tensor issues and modifies `save_pretrained`
        behavior.
    - Initializes the processor if specified as a path or `None`.
    - Sets the minimum tokens per module if `dataset_args` are provided.
    Raises:
        FileNotFoundError: If the model or processor path is invalid.
    """

    # Initialize model
    if isinstance(model_args.model, (str, PosixPath)):
        model = initialize_model_from_path(model_args)
        model_args.model = model

    # Initialize processor if dataset provided
    if isinstance(model_args.processor, (str, type(None))):
        try:
            model_args.processor = initialize_processor_from_path(
                model_args, model_args.model
            )
        except Exception as e:
            if dataset_args.is_dataset_provided():
                raise RuntimeError(
                    "An error occurred when attempting to initialize "
                    "model processor, which is required when a dataset "
                    "is provided. To resolve, create and pass in a "
                    "processor directly to `oneshot`/`train`."
                ) from e
            elif output_dir:
                logger.warning(
                    "Model processor could not be auto-initialized and "
                    "will not be saved along with the model. To resolve, "
                    "create and pass in a processor directly to "
                    f"`oneshot`/`train`.\nInitialization Error: {e}"
                )

    # untie tie_word_embeddings weights
    if not model_args.tie_word_embeddings:
        untie_word_embeddings(model_args.model)

    # if the model was loaded with accelerate offloading, convert to CT offloading
    if hasattr(model_args.model, "hf_device_map"):
        from_accelerate(model_args.model)

    # wrap model.save_pretrained
    modify_save_pretrained(model_args.model)