450 lines
17 KiB
Python
450 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py
|
|
|
|
import itertools
|
|
import re
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
|
|
ROOT_DIR_RELATIVE = '../../../..'
|
|
EXAMPLE_DIR = ROOT_DIR / "examples"
|
|
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
|
|
|
|
|
|
def fix_case(text: str) -> str:
|
|
subs = {
|
|
"api": "API",
|
|
"cli": "CLI",
|
|
"cpu": "CPU",
|
|
"llm": "LLM",
|
|
"tpu": "TPU",
|
|
"aqlm": "AQLM",
|
|
"gguf": "GGUF",
|
|
"lora": "LoRA",
|
|
"rlhf": "RLHF",
|
|
"vllm": "vLLM",
|
|
"openai": "OpenAI",
|
|
"multilora": "MultiLoRA",
|
|
"mlpspeculator": "MLPSpeculator",
|
|
"finetune": "Finetune",
|
|
"distillation": "Distillation",
|
|
"wan": "Wan",
|
|
"i2v": "I2V",
|
|
"t2v": "T2V",
|
|
"1.3b": "1.3B",
|
|
"14b": "14B",
|
|
"480p": "480P",
|
|
"720p": "720P",
|
|
r"fp\d+": lambda x: x.group(0).upper(), # e.g. fp16, fp32
|
|
r"int\d+": lambda x: x.group(0).upper(), # e.g. int8, int16
|
|
}
|
|
for pattern, repl in subs.items():
|
|
text = re.sub(rf'\b{pattern}\b', repl, text,
|
|
flags=re.IGNORECASE) # type: ignore[call-overload]
|
|
return text
|
|
|
|
|
|
@dataclass
|
|
class Index:
|
|
"""
|
|
Index class to generate a structured document index.
|
|
|
|
Attributes:
|
|
path (Path): The path save the index file to.
|
|
title (str): The title of the index.
|
|
description (str): A brief description of the index.
|
|
caption (str): An optional caption for the table of contents.
|
|
maxdepth (int): The maximum depth of the table of contents. Defaults to 1.
|
|
documents (list[str]): A list of document paths to include in the index. Defaults to an empty list.
|
|
|
|
Methods:
|
|
generate() -> str:
|
|
Generates the index content as a string in the specified format.
|
|
""" # noqa: E501
|
|
path: Path
|
|
title: str
|
|
description: str
|
|
caption: str
|
|
maxdepth: int = 1
|
|
documents: list[str] = field(default_factory=list)
|
|
|
|
def generate(self) -> str:
|
|
content = f"# {self.title}\n\n{self.description}\n\n"
|
|
content += ":::{toctree}\n"
|
|
content += f":caption: {self.caption}\n:maxdepth: {self.maxdepth}\n"
|
|
content += "\n".join(self.documents) + "\n:::\n"
|
|
return content
|
|
|
|
|
|
@dataclass
|
|
class Example:
|
|
"""
|
|
Example class for generating documentation content from a given path.
|
|
|
|
Attributes:
|
|
path (Path): The path to the main directory or file.
|
|
category (str): The category of the document.
|
|
main_file (Path): The main file in the directory.
|
|
other_files (list[Path]): list of other files in the directory.
|
|
title (str): The title of the document.
|
|
|
|
Methods:
|
|
__post_init__(): Initializes the main_file, other_files, and title attributes.
|
|
determine_main_file() -> Path: Determines the main file in the given path.
|
|
determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
|
|
determine_title() -> str: Determines the title of the document.
|
|
generate() -> str: Generates the documentation content.
|
|
""" # noqa: E501
|
|
path: Path
|
|
category: str | None = None
|
|
main_file: Path = field(init=False)
|
|
other_files: list[Path] = field(init=False)
|
|
title: str = field(init=False)
|
|
|
|
def __post_init__(self):
|
|
self.main_file = self.determine_main_file()
|
|
self.other_files = self.determine_other_files()
|
|
self.title = self.determine_title()
|
|
|
|
def determine_main_file(self) -> Path:
|
|
"""
|
|
Determines the main file in the given path.
|
|
If the path is a file, it returns the path itself. Otherwise, it searches
|
|
for Markdown files (*.md) in the directory and returns the first one found.
|
|
Returns:
|
|
Path: The main file path, either the original path if it's a file or the first
|
|
Markdown file found in the directory.
|
|
Raises:
|
|
IndexError: If no Markdown files are found in the directory.
|
|
""" # noqa: E501
|
|
return self.path if self.path.is_file() else list(
|
|
self.path.glob("*.md")).pop()
|
|
|
|
def determine_other_files(self) -> list[Path]:
|
|
"""
|
|
Determine other files in the directory excluding the main file.
|
|
|
|
This method checks if the given path is a file. If it is, it returns an empty list.
|
|
Otherwise, it recursively searches through the directory and returns a list of all
|
|
files that are not the main file.
|
|
|
|
Returns:
|
|
list[Path]: A list of Path objects representing the other files in the directory.
|
|
""" # noqa: E501
|
|
if self.path.is_file():
|
|
return []
|
|
is_other_file = lambda file: file.is_file() and file != self.main_file
|
|
return [file for file in self.path.rglob("*")
|
|
if is_other_file(file)] # type: ignore[no-untyped-call]
|
|
|
|
def determine_title(self) -> str:
|
|
return fix_case(self.path.stem.replace("_", " ").title())
|
|
|
|
def generate(self) -> str:
|
|
# Convert the path to a relative path from __file__
|
|
make_relative = lambda path: ROOT_DIR_RELATIVE / path.relative_to(
|
|
ROOT_DIR)
|
|
|
|
content = f"Source <gh-file:{self.path.relative_to(ROOT_DIR)}>.\n\n"
|
|
include = "include" if self.main_file.suffix == ".md" else \
|
|
"literalinclude"
|
|
if include == "literalinclude":
|
|
content += f"# {self.title}\n\n"
|
|
content += f":::{{{include}}} {make_relative(self.main_file)}\n" # type: ignore[no-untyped-call]
|
|
if include == "literalinclude":
|
|
content += f":language: {self.main_file.suffix[1:]}\n"
|
|
content += ":::\n\n"
|
|
|
|
if not self.other_files:
|
|
return content
|
|
|
|
content += "## Example materials\n\n"
|
|
for file in sorted(self.other_files):
|
|
include = "include" if file.suffix == ".md" else "literalinclude"
|
|
content += f":::{{admonition}} {file.relative_to(self.path)}\n"
|
|
content += ":class: dropdown\n\n"
|
|
content += f":::{{{include}}} {make_relative(file)}\n:::\n" # type: ignore[no-untyped-call]
|
|
content += ":::\n\n"
|
|
|
|
return content
|
|
|
|
|
|
@dataclass
|
|
class NestedStructure:
|
|
"""Helper class to manage nested documentation structures for training/distillation."""
|
|
category: str
|
|
method: str
|
|
model: str
|
|
dataset: str
|
|
example: Example
|
|
|
|
@property
|
|
def filename(self) -> str:
|
|
return f"{self.model}_{self.dataset}"
|
|
|
|
@property
|
|
def title(self) -> str:
|
|
return fix_case(self.dataset.replace('_', ' '))
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
category_name = self.category.title()
|
|
return f"{category_name} example using the {self.dataset} dataset with the {self.model} model."
|
|
|
|
|
|
def create_category_indices() -> dict[str, Index]:
|
|
"""Create category indices with their respective configurations."""
|
|
main_index_dir = ROOT_DIR / "docs/source/examples"
|
|
if not main_index_dir.exists():
|
|
main_index_dir.mkdir(parents=True)
|
|
|
|
category_indices = {
|
|
"inference":
|
|
Index(
|
|
path=ROOT_DIR /
|
|
"docs/source/inference/examples/examples_inference_index.md",
|
|
title="🚀 Examples",
|
|
description=
|
|
"Inference examples demonstrate how to use FastVideo inference. We recommend starting with <project:basic.md>.",
|
|
caption="Examples",
|
|
maxdepth=1,
|
|
),
|
|
"training":
|
|
Index(
|
|
path=ROOT_DIR /
|
|
"docs/source/training/examples/examples_training_index.md",
|
|
title="🚀 Examples",
|
|
description=
|
|
"Training examples demonstrate how to use FastVideo training.",
|
|
caption="Examples",
|
|
maxdepth=3,
|
|
),
|
|
"distillation":
|
|
Index(
|
|
path=ROOT_DIR /
|
|
"docs/source/distillation/examples/examples_distillation_index.md",
|
|
title="🚀 Examples",
|
|
description=
|
|
"Distillation examples demonstrate how to use FastVideo distillation.",
|
|
caption="Examples",
|
|
maxdepth=3,
|
|
),
|
|
}
|
|
|
|
# Ensure all category doc directories exist
|
|
for index in category_indices.values():
|
|
if not index.path.parent.exists():
|
|
index.path.parent.mkdir(parents=True)
|
|
|
|
return category_indices
|
|
|
|
|
|
def find_examples(category_indices: dict[str, Index],
|
|
generate_main_index: bool) -> list[Example]:
|
|
"""Find all examples from the examples directory."""
|
|
examples = []
|
|
glob_patterns = ["*.py", "*.md", "*.sh"]
|
|
|
|
# Find categorised examples
|
|
for category in category_indices:
|
|
category_dir = EXAMPLE_DIR / category
|
|
globs = [category_dir.glob(pattern) for pattern in glob_patterns]
|
|
for path in itertools.chain(*globs):
|
|
examples.append(Example(path, category))
|
|
# Find examples in subdirectories (recursively)
|
|
for path in category_dir.glob("**/*.md"):
|
|
examples.append(Example(path.parent, category))
|
|
|
|
# Find uncategorised examples only if we're generating a main index
|
|
if generate_main_index:
|
|
globs = [EXAMPLE_DIR.glob(pattern) for pattern in glob_patterns]
|
|
for path in itertools.chain(*globs):
|
|
examples.append(Example(path))
|
|
# Find examples in subdirectories
|
|
for path in EXAMPLE_DIR.glob("*/*.md"):
|
|
# Skip categorised examples
|
|
if path.parent.name in category_indices:
|
|
continue
|
|
examples.append(Example(path.parent))
|
|
|
|
return examples
|
|
|
|
|
|
def create_nested_structures(
|
|
examples: list[Example]
|
|
) -> dict[str, dict[str, dict[str, dict[str, NestedStructure]]]]:
|
|
"""Create nested structures for training and distillation categories."""
|
|
nested_structures: dict[str, dict[str, dict[str,
|
|
dict[str,
|
|
NestedStructure]]]] = {}
|
|
|
|
for example in examples:
|
|
if example.category not in ["training", "distillation"]:
|
|
continue
|
|
|
|
category_dir = EXAMPLE_DIR / example.category
|
|
relative_path = example.path.relative_to(category_dir)
|
|
path_parts = relative_path.parts
|
|
|
|
# For nested examples like finetune/wan_i2v_14b_480p/crush_smol
|
|
if len(path_parts) >= 3:
|
|
method = path_parts[0] # e.g., "finetune"
|
|
model = path_parts[1] # e.g., "wan_i2v_14b_480p"
|
|
dataset = path_parts[2] # e.g., "crush_smol"
|
|
|
|
# Initialize nested structure
|
|
if example.category not in nested_structures:
|
|
nested_structures[example.category] = {}
|
|
if method not in nested_structures[example.category]:
|
|
nested_structures[example.category][method] = {}
|
|
if model not in nested_structures[example.category][method]:
|
|
nested_structures[example.category][method][model] = {}
|
|
|
|
# Store the nested structure
|
|
nested_structures[
|
|
example.category][method][model][dataset] = NestedStructure(
|
|
category=example.category,
|
|
method=method,
|
|
model=model,
|
|
dataset=dataset,
|
|
example=example)
|
|
|
|
return nested_structures
|
|
|
|
|
|
def generate_flat_examples(examples: list[Example],
|
|
category_indices: dict[str, Index],
|
|
examples_index: Index | None,
|
|
generate_main_index: bool) -> None:
|
|
"""Generate documentation for flat structure examples (inference, etc.)."""
|
|
for example in examples:
|
|
if example.category in ["training", "distillation"]:
|
|
continue # Skip nested structure examples
|
|
|
|
# Determine which index to use for this example
|
|
if example.category is not None and example.category in category_indices:
|
|
index = category_indices[example.category]
|
|
elif generate_main_index:
|
|
assert examples_index is not None
|
|
index = examples_index
|
|
else:
|
|
continue
|
|
|
|
# Generate the example documentation
|
|
doc_path = index.path.parent / f"{example.path.stem}.md"
|
|
with open(doc_path, "w+") as f:
|
|
f.write(example.generate())
|
|
index.documents.append(example.path.stem)
|
|
|
|
|
|
def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
|
|
str, dict[str, NestedStructure]]]], category_indices: dict[str,
|
|
Index]) -> None:
|
|
"""Generate documentation for nested structure examples (training, distillation)."""
|
|
for category_name in ["training", "distillation"]:
|
|
if category_name not in category_indices or category_name not in nested_structures:
|
|
continue
|
|
|
|
category_index = category_indices[category_name]
|
|
category_base_dir = category_index.path.parent
|
|
|
|
for method, models in nested_structures[category_name].items():
|
|
# Create method-level index
|
|
method_index = Index(path=category_base_dir / f"{method}.md",
|
|
title=fix_case(method),
|
|
description=f"Examples using {method}.",
|
|
caption=f"{fix_case(method)} Examples",
|
|
maxdepth=2)
|
|
|
|
for model, datasets in models.items():
|
|
# Generate dataset examples using the Example class
|
|
for dataset, nested_struct in datasets.items():
|
|
doc_path = category_base_dir / f"{nested_struct.filename}.md"
|
|
with open(doc_path, "w+") as f:
|
|
f.write(nested_struct.example.generate())
|
|
|
|
# Create model-level index
|
|
model_index = Index(
|
|
path=category_base_dir / f"{model}.md",
|
|
title=fix_case(model.replace('_', ' ')),
|
|
description=f"Examples for the {model} model.",
|
|
caption=f"{fix_case(model.replace('_', ' '))} Datasets",
|
|
maxdepth=1)
|
|
|
|
# Add dataset indices to model index
|
|
for dataset, nested_struct in datasets.items():
|
|
model_index.documents.append(nested_struct.filename)
|
|
|
|
# Write model index
|
|
with open(model_index.path, "w+") as f:
|
|
f.write(model_index.generate())
|
|
|
|
# Add model to method index
|
|
method_index.documents.append(model)
|
|
|
|
# Write method index
|
|
with open(method_index.path, "w+") as f:
|
|
f.write(method_index.generate())
|
|
|
|
# Add method to main category index
|
|
category_index.documents.append(method)
|
|
|
|
|
|
def generate_examples(generate_main_index=False):
|
|
"""
|
|
Generate example documentation.
|
|
|
|
Args:
|
|
generate_main_index (bool): Whether to generate the main examples index.
|
|
If False, only category-specific indices will be generated.
|
|
"""
|
|
# Create category indices
|
|
category_indices = create_category_indices()
|
|
|
|
# Create the main examples index only if requested
|
|
examples_index = None
|
|
if generate_main_index:
|
|
main_index_dir = ROOT_DIR / "docs/source/examples"
|
|
examples_index = Index(
|
|
path=main_index_dir / "examples_index.md",
|
|
title="💡 Examples",
|
|
description=
|
|
"A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",
|
|
caption="Examples",
|
|
maxdepth=2)
|
|
|
|
# Find all examples
|
|
examples = find_examples(category_indices, generate_main_index)
|
|
|
|
# Create nested structures for training and distillation
|
|
nested_structures = create_nested_structures(examples)
|
|
|
|
# Generate flat structure examples (inference, etc.)
|
|
generate_flat_examples(examples, category_indices, examples_index,
|
|
generate_main_index)
|
|
|
|
# Generate nested structure examples (training, distillation)
|
|
generate_nested_examples(nested_structures, category_indices)
|
|
|
|
# Generate the index files for categories
|
|
for category_index in category_indices.values():
|
|
if category_index.documents:
|
|
# Add to main index if it exists
|
|
if generate_main_index and examples_index:
|
|
main_index_dir = examples_index.path.parent
|
|
rel_path = category_index.path.relative_to(
|
|
main_index_dir.parent)
|
|
examples_index.documents.insert(
|
|
0,
|
|
str(rel_path).replace(".md", ""))
|
|
|
|
# Write the category index file
|
|
with open(category_index.path, "w+") as f:
|
|
f.write(category_index.generate())
|
|
|
|
# Write the main index file if requested
|
|
if generate_main_index and examples_index:
|
|
with open(examples_index.path, "w+") as f:
|
|
f.write(examples_index.generate())
|