llmcompressor.core.session

def event(
    self,
    event_type: EventType,
    batch_data: Optional[Any] = None,
    loss: Optional[Any] = None,
    **kwargs,
) -> ModifiedState:
    """
    Invoke an event for current CompressionSession.

    :param event_type: the event type to invoke
    :param batch_data: the batch data to use for the event
    :param loss: the loss to use for the event if any
    :param kwargs: additional kwargs to pass to the lifecycle's event method
    :return: the modified state of the session after invoking the event
    """
    mod_data = self._lifecycle.event(
        event_type=event_type, batch_data=batch_data, loss=loss, **kwargs
    )
    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,  # TODO: is this supposed to be a different type?
        modifier_data=mod_data,
    )

def finalize(self, **kwargs) -> ModifiedState:
    """
    Finalize the session for compression. This will run the finalize method
    for each modifier in the session's lifecycle. This will also set the session's
    state to the finalized state.

    :param kwargs: additional kwargs to pass to the lifecycle's finalize method
    :return: the modified state of the session after finalizing
    """
    mod_data = self._lifecycle.finalize(**kwargs)

    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,
        modifier_data=mod_data,
    )

def get_serialized_recipe(self) -> Optional[str]:
    """
    :return: serialized string of the current compiled recipe
    """
    recipe = self.lifecycle.recipe

    if recipe is not None and hasattr(recipe, "yaml"):
        return recipe.yaml()

    logger.warning("Recipe not found in session - it may have been reset")

def initialize(
    self,
    recipe: Union[str, List[str], "Recipe", List["Recipe"], None] = None,
    recipe_stage: Union[str, List[str], None] = None,
    recipe_args: Union[Dict[str, Any], None] = None,
    model: Optional[Any] = None,
    teacher_model: Optional[Any] = None,
    optimizer: Optional[Any] = None,
    attach_optim_callbacks: bool = True,
    train_data: Optional[Any] = None,
    val_data: Optional[Any] = None,
    test_data: Optional[Any] = None,
    calib_data: Optional[Any] = None,
    copy_data: bool = True,
    start: Optional[float] = None,
    steps_per_epoch: Optional[int] = None,
    batches_per_step: Optional[int] = None,
    loggers: Union[None, LoggerManager, List[BaseLogger]] = None,
    **kwargs,
) -> ModifiedState:
    """
    Initialize the session for compression. This will run the initialize method
    for each modifier in the session's lifecycle. This will also set the session's
    state to the initialized state.

    :param recipe: the recipe to use for the compression, can be a path to a
        recipe file, a raw recipe string, a recipe object, or a list
        of recipe objects.
    :param recipe_stage: the stage to target for the compression
    :param recipe_args: the args to use for overriding the recipe defaults
    :param model: the model to compress
    :param teacher_model: the teacher model to use for knowledge distillation
    :param optimizer: the optimizer to use for the compression
    :param attach_optim_callbacks: True to attach the optimizer callbacks to the
        compression lifecycle, False otherwise
    :param train_data: the training data to use for the compression
    :param val_data: the validation data to use for the compression
    :param test_data: the testing data to use for the compression
    :param calib_data: the calibration data to use for the compression
    :param copy_data: True to copy the data, False otherwise
    :param start: the start epoch to use for the compression
    :param steps_per_epoch: the number of steps per epoch to use for the
        compression
    :param batches_per_step: the number of batches per step to use for
        compression
    :param loggers: the metrics manager to setup logging important info
        and milestones to, also accepts a list of BaseLogger(s)
    :param kwargs: additional kwargs to pass to the lifecycle's initialize method
    :return: the modified state of the session after initializing
    """
    mod_data = self._lifecycle.initialize(
        recipe=recipe,
        recipe_stage=recipe_stage,
        recipe_args=recipe_args,
        model=model,
        teacher_model=teacher_model,
        optimizer=optimizer,
        attach_optim_callbacks=attach_optim_callbacks,
        train_data=train_data,
        val_data=val_data,
        test_data=test_data,
        calib_data=calib_data,
        copy_data=copy_data,
        start=start,
        steps_per_epoch=steps_per_epoch,
        batches_per_step=batches_per_step,
        loggers=loggers,
        **kwargs,
    )

    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,
        modifier_data=mod_data,
    )

def log(self, event_type: EventType, loss: Optional[Any] = None):
    """
    Log model and loss information for the current event type

    :param event_type: the event type to log for
    :param loss: the loss to log if any
    """
    self._log_model_info()
    self._log_loss(event_type=event_type, loss=loss)

def reset(self):
    """
    Reset the session to its initial state
    """
    self._lifecycle.reset()

def reset_stage(self):
    """
    Reset the session for starting a new stage, recipe and model stays intact
    """
    self.lifecycle.initialized_ = False
    self.lifecycle.finalized = False