Updating API docs

Signed-off-by: Aidan Reilly <[email protected]>

Author: aireilly

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,19 @@ plugins:
       handlers:
         python:
           options:
-            docstring_style: sphinx
+            docstring_section_style: list
+            show_if_no_docstring: false
+            docstring_style: "sphinx"
+            filters: ["!^_"]
+            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