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,18 @@ plugins:
       handlers:
         python:
           options:
-            docstring_style: sphinx
+            docstring_section_style: list
+            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,6 @@
-"""
-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,11 @@
 # 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. Support 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.
+
+This module 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.
+
+This module 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,10 @@
+""" Helper functions for core compression operations.
+
+This module provides utility functions for logging model information and state
+    management during compression workflows. Includes functionality for
+    conditional logging, model introspection, and parameter tracking.
+"""
+
 from typing import Any, Generator, Optional, Tuple, Union
 
 from llmcompressor.core.state import State