Updating API docs

mkdocs.yml

@@ -70,6 +70,7 @@ markdown_extensions:
   - pymdownx.mark
   - pymdownx.smartsymbols
   - pymdownx.snippets
+  - pymdownx.magiclink
   - pymdownx.superfences:
       custom_fences:
         - name: mermaid
@@ -85,6 +86,9 @@ markdown_extensions:
 plugins:
   - api-autonav:
       modules: ['src/llmcompressor']
+      exclude:
+        - "llmcompressor.version"
+        - "llmcompressor.typing"
   - gen-files:
       scripts:
         - docs/scripts/gen_files.py
@@ -96,7 +100,21 @@ plugins:
       handlers:
         python:
           options:
-            docstring_style: sphinx
+            docstring_section_style: list
+            show_if_no_docstring: false
+            docstring_style: "sphinx"
+            filters:
+              - "!docstring"
+              - "!^_"
+            heading_level: 1
+            merge_init_into_class: true
+            parameter_headings: true
+            separate_signature: true
+            show_root_heading: true
+            show_signature_annotations: true
+            show_symbol_type_heading: true
+            show_symbol_type_toc: true
+            summary: true
   - search
   - section-index
   - social

src/llmcompressor/__init__.py

@@ -1,6 +1,7 @@
 """
-A library for compressing large language models utilizing the latest techniques and
-research in the field for both training aware and post training techniques.
+LLM Compressor is a library for compressing large language models utilizing
+the latest techniques and research in the field for both training aware and
+post-training techniques.
 
 The library is designed to be flexible and easy to use on top of
 PyTorch and HuggingFace Transformers, allowing for quick experimentation.

src/llmcompressor/args/__init__.py

@@ -1,5 +1,12 @@
 # flake8: noqa
 
+"""
+Arguments package for LLM Compressor.
+
+Defines structured argument classes for datasets, models, training, and
+recipes, along with utilities for parsing them.
+"""
+
 from .dataset_arguments import DatasetArguments
 from .model_arguments import ModelArguments
 from .recipe_arguments import RecipeArguments

src/llmcompressor/args/dataset_arguments.py

@@ -1,3 +1,12 @@
+"""
+Dataset argument classes for LLM compression workflows.
+
+This module defines dataclass-based argument containers for configuring dataset
+loading, preprocessing, and calibration parameters across different dataset
+sources and processing pipelines. Supports various input formats including
+HuggingFace datasets, custom JSON/CSV files, and DVC-managed datasets.
+"""
+
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, List, Optional, Union
 

src/llmcompressor/args/model_arguments.py

@@ -1,3 +1,12 @@
+"""
+Model argument classes for LLM compression workflows.
+
+This module defines dataclass-based argument containers for configuring model
+loading, tokenization, and preprocessing parameters. Supports various model
+sources including HuggingFace model hub, local paths, and custom
+configurations for compression workflows.
+"""
+
 from dataclasses import dataclass, field
 from typing import Optional
 

src/llmcompressor/args/recipe_arguments.py

@@ -1,3 +1,11 @@
+"""
+Recipe argument classes for LLM compression workflows.
+
+Defines dataclass-based argument containers for configuring sparsification
+recipes, compression sessions, and stage-based execution parameters used in
+model compression and optimization workflows.
+"""
+
 from dataclasses import dataclass, field
 from typing import List, Optional
 

src/llmcompressor/args/training_arguments.py

@@ -1,3 +1,12 @@
+"""
+Training argument classes for LLM compression workflows.
+
+This module defines dataclass-based argument containers for configuring
+training and one-shot calibration workflows. Extends HuggingFace's
+TrainingArguments with additional parameters specific to compression and
+stage-based execution.
+"""
+
 from dataclasses import dataclass, field
 from typing import Optional
 

src/llmcompressor/args/utils.py

@@ -1,3 +1,12 @@
+"""
+Utility functions for parsing and processing argument classes.
+
+Provides helper functions for parsing command-line arguments and
+configuration dictionaries into structured argument dataclasses used in
+LLM compression workflows. Handles argument validation, deprecation
+warnings, and processor resolution.
+"""
+
 from typing import Tuple
 
 from loguru import logger

src/llmcompressor/core/__init__.py

@@ -1,3 +1,11 @@
+"""
+Provides the core compression framework for LLM Compressor.
+
+The core API manages compression sessions, tracks state changes, handles events
+during compression, and Provides lifecycle hooks for the compression
+process.
+"""
+
 from llmcompressor.core.events import Event, EventType
 from llmcompressor.core.lifecycle import CompressionLifecycle
 from llmcompressor.core.model_layer import ModelParameterizedLayer

src/llmcompressor/core/helpers.py

@@ -1,3 +1,11 @@
+"""
+Helper functions for core compression operations.
+
+Provides utility functions for logging model information and state
+management during compression workflows. Includes functionality for
+conditional logging and parameter tracking.
+"""
+
 from typing import Any, Generator, Optional, Tuple, Union
 
 from llmcompressor.core.state import State

src/llmcompressor/core/lifecycle.py

@@ -1,8 +1,8 @@
 """
 Module for managing the compression lifecycle in the LLM Compressor.
 
-This module provides a class for defining and managing the lifecycle of compression
-events, including initialization, finalization, and event handling.
+Provides a class for defining and managing the lifecycle of compression events,
+including initialization, finalization, and event handling.
 """
 
 from dataclasses import dataclass, field

src/llmcompressor/core/model_layer.py

@@ -1,3 +1,11 @@
+"""
+Model layer utility classes for LLM compression workflows.
+
+Provides dataclass containers for managing model layers and their associated
+parameters during compression operations. Facilitates tracking and manipulation
+of specific model components and their parameters.
+"""
+
 from dataclasses import dataclass
 from typing import Any
 

src/llmcompressor/core/session.py

@@ -1,3 +1,11 @@
+"""
+Compression session management for LLM compression workflows.
+
+Provides the main CompressionSession class for managing compression
+workflows, including lifecycle management, event handling, callback
+registration, and state tracking.
+"""
+
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -47,7 +55,7 @@ def __init__(self):
     def lifecycle(self) -> CompressionLifecycle:
         """
         Lifecycle is used to keep track of where we are in the compression
-        process and what modifiers are active. It also provides the ability
+        process and what modifiers are active. It also Provides the ability
         to invoke events on the lifecycle.
 
         :return: the lifecycle for the session

src/llmcompressor/core/session_functions.py

@@ -1,3 +1,10 @@
+"""
+Session management functions for LLM compression workflows.
+
+Provides global session management functionality including session
+creation, activation, reset operations, and lifecycle callback management.
+"""
+
 import threading
 from contextlib import contextmanager
 from typing import Any, Generator, Optional

src/llmcompressor/core/state.py

@@ -1,7 +1,7 @@
 """
-Module for managing the state of the LLM Compressor.
+Module for managing LLM Compressor state.
 
-This module provides classes for holding and updating the state information
+Provides classes for holding and updating the state information
 related to data, hardware, and model compression.
 """
 

src/llmcompressor/datasets/__init__.py

@@ -1,5 +1,12 @@
 # flake8: noqa
 
+"""
+Provides dataset utilities for model calibration and processing.
+
+Includes functions to format calibration data, create dataloaders,
+process datasets, and split datasets for quantization workflows.
+"""
+
 from .utils import (
     format_calibration_data,
     get_calibration_dataloader,

src/llmcompressor/datasets/utils.py

@@ -1,3 +1,12 @@
+"""
+Dataset utility functions for LLM compression workflows.
+
+Provides helper functions for loading, processing, and formatting datasets used
+in model compression pipelines. Handles dataset splitting, tokenization,
+calibration data preparation, and dataloader creation for both training and
+one-shot calibration workflows.
+"""
+
 import re
 from typing import Any, Callable, Dict, List, Optional
 

src/llmcompressor/entrypoints/__init__.py

@@ -1,4 +1,12 @@
 # flake8: noqa
+
+"""
+Provides entry points for model compression workflows.
+
+Includes oneshot compression, training, and pre and post-processing utilities
+for model optimization tasks.
+"""
+
 from .oneshot import Oneshot, oneshot
 from .train import train
 from .utils import post_process, pre_process

src/llmcompressor/entrypoints/oneshot.py

@@ -1,3 +1,12 @@
+"""
+Oneshot compression entrypoint for post-training model optimization.
+
+Provides the main oneshot compression entry point for applying
+quantization, pruning, and other compression techniques to pre-trained
+models without additional training. Supports calibration-based compression
+with various pipeline configurations for efficient model optimization.
+"""
+
 import os
 from datetime import datetime
 from typing import TYPE_CHECKING, List, Optional, Union

src/llmcompressor/entrypoints/train.py

@@ -1,3 +1,12 @@
+"""
+Training entrypoint for fine-tuning models with compression support.
+
+Provides the main training entry point that supports both vanilla
+fine-tuning and compression-aware training workflows. Integrates with
+HuggingFace transformers and supports knowledge distillation, pruning,
+and quantization during the training process.
+"""
+
 import math
 import os
 

src/llmcompressor/entrypoints/utils.py

@@ -1,3 +1,12 @@
+"""
+Utility functions for entrypoint pre and post-processing operations.
+
+Provides common utility functions used by training and one-shot
+compression entrypoints. Includes model loading, configuration setup,
+preprocessing steps, and post-processing operations for compression
+workflows.
+"""
+
 import inspect
 import os
 from pathlib import PosixPath

src/llmcompressor/logger.py

@@ -1,21 +1,24 @@
 """
-Logger configuration for LLM Compressor.
-
-This module provides a flexible logging configuration using the loguru library.
-It supports console and file logging with options to configure via environment
-variables or direct function calls.
-
-Environment Variables:
-    - LLM_COMPRESSOR_LOG_DISABLED: Disable logging (default: false).
-    - LLM_COMPRESSOR_CLEAR_LOGGERS: Clear existing loggers from loguru (default: true).
-    - LLM_COMPRESSOR_LOG_LEVEL: Log level for console logging
-        (default: none, options: DEBUG, INFO, WARNING, ERROR, CRITICAL).
-    - LLM_COMPRESSOR_LOG_FILE: Path to the log file for file logging
-        (default: llm-compressor.log if log file level set else none)
-    - LLM_COMPRESSOR_LOG_FILE_LEVEL: Log level for file logging
-        (default: INFO if log file set else none).
-
-Usage:
+Provides a flexible logging configuration for LLM Compressor.
+
+Using the loguru library, Logger supports console and file logging with options
+to configure via environment variables or direct function calls.
+
+**Environment Variables**
+
+- `LLM_COMPRESSOR_LOG_DISABLED`: Disable logging.
+    Default: `False`.
+- `LLM_COMPRESSOR_CLEAR_LOGGERS`: Clear existing loggers from loguru.
+    Default: `True`.
+- `LLM_COMPRESSOR_LOG_LEVEL`: Log level for console logging.
+    Default: `None`. Options: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL`.
+- `LLM_COMPRESSOR_LOG_FILE`: Path to the log file for file logging.
+    Default: `llm-compressor.log` if log file level set, otherwise `None`.
+- `LLM_COMPRESSOR_LOG_FILE_LEVEL`: Log level for file logging.
+    Default: `INFO` if log file is set, otherwise `None`.
+
+**Usage**
+
     from llmcompressor import logger, configure_logger, LoggerConfig
 
     # Configure metrics with default settings
@@ -60,10 +63,11 @@ class LoggerConfig:
 def configure_logger(config: Optional[LoggerConfig] = None) -> None:
     """
     Configure the logger for LLM Compressor.
+
     This function sets up the console and file logging
     as per the specified or default parameters.
 
-    Note: Environment variables take precedence over the function parameters.
+    **Note**: Environment variables take precedence over function parameters.
 
     :param config: The configuration for the logger to use.
     :type config: LoggerConfig

src/llmcompressor/metrics/__init__.py

@@ -1,3 +1,12 @@
+"""
+Metrics logging and monitoring framework for compression workflows.
+
+Provides comprehensive metrics collection, logging, and monitoring
+capabilities for model compression operations. Includes base loggers,
+frequency management, and specialized metrics tracking for training and
+inference performance during compression.
+"""
+
 # flake8: noqa
 
 from .logger import *

src/llmcompressor/metrics/logger.py

@@ -1,5 +1,6 @@
 """
-Contains code for loggers that help visualize the information from each modifier
+Contains code for loggers that help visualize the information from each
+modifier.
 """
 
 import os

src/llmcompressor/metrics/utils/frequency_manager.py

@@ -1,3 +1,12 @@
+"""
+Frequency management utilities for metrics logging.
+
+Provides classes and functions for managing logging frequencies
+and determining when metrics should be recorded during training and
+compression workflows. Supports both epoch-based and step-based logging
+with configurable modes and intervals.
+"""
+
 from typing import Literal, Optional, Union
 
 __all__ = [

src/llmcompressor/modeling/__init__.py

@@ -1,3 +1,12 @@
+"""
+Model preparation and fusion utilities for compression workflows.
+
+Provides tools for preparing models for compression including
+layer fusion, module preparation, and model structure optimization.
+Handles pre-compression transformations and architectural modifications
+needed for efficient compression.
+"""
+
 # flake8: noqa
 
 from .fuse import *