llmcompressor.entrypoints.utils

def get_processor_name_from_model(student: Module, teacher: Optional[Module]) -> str:
    """
    Get a processor/tokenizer source used for both student and teacher, assuming
    that they could be shared

    :param student: the student model
    :param teacher: the teacher model
    :return: the source for the processor/tokenizer shared between teacher and model
    """
    if teacher is not None and teacher not in ("disable", "self"):
        student_forward_params = list(
            inspect.signature(student.forward).parameters.keys()
        )
        teacher_forward_params = list(
            inspect.signature(teacher.forward).parameters.keys()
        )
        diff = [p for p in student_forward_params if p not in teacher_forward_params]
        if diff:
            raise RuntimeError(
                "Teacher tokenizer cannot be used for student "
                f"due to missing args: {diff}"
            )
        src_model = teacher
    else:
        src_model = student
    return src_model.config._name_or_path

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

    Save is skipped for stage runs for `train` - saves using the trainer.save_model()

    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.
    """
    # remove any existing dispatches
    if model_args is not None and model_args.model is not None:
        remove_dispatch(model_args.model)

    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"):
    """
    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.
    """
    _warn_tied_embeddings(model_args.tie_word_embeddings)

    # Initialize model
    if isinstance(model_args.model, (str, PosixPath)):
        model, distill_teacher = initialize_model_from_path(model_args)
        if is_fsdp_model(model):
            raise NotImplementedError(
                "FSDP models are not supported in the current release but will be "
                "suported in future releases of LLM Compressor."
            )
        model_args.model = model
        model_args.distill_teacher = distill_teacher

    # Initialize processor
    if isinstance(model_args.processor, (str, type(None))):
        model_args.processor = initialize_processor_from_path(
            model_args, model_args.model
        )

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

    # wrap model.save_pretrained
    modify_save_pretrained(model_args.model)