|
| 1 | +""" Sourced from https://github.com/vllm-project/vllm/blob/main/docs/source/generate_examples.py """ # noqa: E501 |
| 2 | + |
| 3 | +import itertools |
| 4 | +import re |
| 5 | +from dataclasses import dataclass, field |
| 6 | +from pathlib import Path |
| 7 | + |
| 8 | +ROOT_DIR = Path(__file__).parent.parent.parent.resolve() |
| 9 | +ROOT_DIR_RELATIVE = '../../../..' |
| 10 | +EXAMPLE_DIR = ROOT_DIR / "examples" |
| 11 | +EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples" |
| 12 | + |
| 13 | + |
| 14 | +def fix_case(text: str) -> str: |
| 15 | + subs = { |
| 16 | + "api": "API", |
| 17 | + "cli": "CLI", |
| 18 | + "cpu": "CPU", |
| 19 | + "llm": "LLM", |
| 20 | + "mae": "MAE", |
| 21 | + "tpu": "TPU", |
| 22 | + "aqlm": "AQLM", |
| 23 | + "gguf": "GGUF", |
| 24 | + "lora": "LoRA", |
| 25 | + "rlhf": "RLHF", |
| 26 | + "vllm": "vLLM", |
| 27 | + "openai": "OpenAI", |
| 28 | + "lmcache": "LMCache", |
| 29 | + "multilora": "MultiLoRA", |
| 30 | + "mlpspeculator": "MLPSpeculator", |
| 31 | + r"fp\d+": lambda x: x.group(0).upper(), # e.g. fp16, fp32 |
| 32 | + r"int\d+": lambda x: x.group(0).upper(), # e.g. int8, int16 |
| 33 | + } |
| 34 | + for pattern, repl in subs.items(): |
| 35 | + text = re.sub(rf'\b{pattern}\b', repl, text, flags=re.IGNORECASE) |
| 36 | + return text |
| 37 | + |
| 38 | + |
| 39 | +@dataclass |
| 40 | +class Index: |
| 41 | + """ |
| 42 | + Index class to generate a structured document index. |
| 43 | +
|
| 44 | + Attributes: |
| 45 | + path (Path): The path save the index file to. |
| 46 | + title (str): The title of the index. |
| 47 | + description (str): A brief description of the index. |
| 48 | + caption (str): An optional caption for the table of contents. |
| 49 | + maxdepth (int): The maximum depth of the table of contents. Defaults to 1. |
| 50 | + documents (list[str]): A list of document paths to include in the index. Defaults to an empty list. |
| 51 | +
|
| 52 | + Methods: |
| 53 | + generate() -> str: |
| 54 | + Generates the index content as a string in the specified format. |
| 55 | + """ # noqa: E501 |
| 56 | + path: Path |
| 57 | + title: str |
| 58 | + description: str |
| 59 | + caption: str |
| 60 | + maxdepth: int = 1 |
| 61 | + documents: list[str] = field(default_factory=list) |
| 62 | + |
| 63 | + def generate(self) -> str: |
| 64 | + content = f"# {self.title}\n\n{self.description}\n\n" |
| 65 | + content += ":::{toctree}\n" |
| 66 | + content += f":caption: {self.caption}\n:maxdepth: {self.maxdepth}\n" |
| 67 | + content += "\n".join(self.documents) + "\n:::\n" |
| 68 | + return content |
| 69 | + |
| 70 | + |
| 71 | +@dataclass |
| 72 | +class Example: |
| 73 | + """ |
| 74 | + Example class for generating documentation content from a given path. |
| 75 | +
|
| 76 | + Attributes: |
| 77 | + path (Path): The path to the main directory or file. |
| 78 | + category (str): The category of the document. |
| 79 | + main_file (Path): The main file in the directory. |
| 80 | + other_files (list[Path]): list of other files in the directory. |
| 81 | + title (str): The title of the document. |
| 82 | +
|
| 83 | + Methods: |
| 84 | + __post_init__(): Initializes the main_file, other_files, and title attributes. |
| 85 | + determine_main_file() -> Path: Determines the main file in the given path. |
| 86 | + determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file. |
| 87 | + determine_title() -> str: Determines the title of the document. |
| 88 | + generate() -> str: Generates the documentation content. |
| 89 | + """ # noqa: E501 |
| 90 | + path: Path |
| 91 | + category: str = None |
| 92 | + main_file: Path = field(init=False) |
| 93 | + other_files: list[Path] = field(init=False) |
| 94 | + title: str = field(init=False) |
| 95 | + |
| 96 | + def __post_init__(self): |
| 97 | + self.main_file = self.determine_main_file() |
| 98 | + self.other_files = self.determine_other_files() |
| 99 | + self.title = self.determine_title() |
| 100 | + |
| 101 | + def determine_main_file(self) -> Path: |
| 102 | + """ |
| 103 | + Determines the main file in the given path. |
| 104 | + If the path is a file, it returns the path itself. Otherwise, it searches |
| 105 | + for Markdown files (*.md) in the directory and returns the first one found. |
| 106 | + Returns: |
| 107 | + Path: The main file path, either the original path if it's a file or the first |
| 108 | + Markdown file found in the directory. |
| 109 | + Raises: |
| 110 | + IndexError: If no Markdown files are found in the directory. |
| 111 | + """ # noqa: E501 |
| 112 | + return self.path if self.path.is_file() else list( |
| 113 | + self.path.glob("*.md")).pop() |
| 114 | + |
| 115 | + def determine_other_files(self) -> list[Path]: |
| 116 | + """ |
| 117 | + Determine other files in the directory excluding the main file. |
| 118 | +
|
| 119 | + This method checks if the given path is a file. If it is, it returns an empty list. |
| 120 | + Otherwise, it recursively searches through the directory and returns a list of all |
| 121 | + files that are not the main file. |
| 122 | +
|
| 123 | + Returns: |
| 124 | + list[Path]: A list of Path objects representing the other files in the directory. |
| 125 | + """ # noqa: E501 |
| 126 | + if self.path.is_file(): |
| 127 | + return [] |
| 128 | + is_other_file = lambda file: file.is_file() and file != self.main_file |
| 129 | + return [file for file in self.path.rglob("*") if is_other_file(file)] |
| 130 | + |
| 131 | + def determine_title(self) -> str: |
| 132 | + return fix_case(self.path.stem.replace("_", " ").title()) |
| 133 | + |
| 134 | + def generate(self) -> str: |
| 135 | + # Convert the path to a relative path from __file__ |
| 136 | + make_relative = lambda path: ROOT_DIR_RELATIVE / path.relative_to( |
| 137 | + ROOT_DIR) |
| 138 | + |
| 139 | + content = f"Source <gh-file:{self.path.relative_to(ROOT_DIR)}>.\n\n" |
| 140 | + include = "include" if self.main_file.suffix == ".md" else \ |
| 141 | + "literalinclude" |
| 142 | + if include == "literalinclude": |
| 143 | + content += f"# {self.title}\n\n" |
| 144 | + content += f":::{{{include}}} {make_relative(self.main_file)}\n" |
| 145 | + if include == "literalinclude": |
| 146 | + content += f":language: {self.main_file.suffix[1:]}\n" |
| 147 | + content += ":::\n\n" |
| 148 | + |
| 149 | + if not self.other_files: |
| 150 | + return content |
| 151 | + |
| 152 | + content += "## Example materials\n\n" |
| 153 | + for file in sorted(self.other_files): |
| 154 | + include = "include" if file.suffix == ".md" else "literalinclude" |
| 155 | + content += f":::{{admonition}} {file.relative_to(self.path)}\n" |
| 156 | + content += ":class: dropdown\n\n" |
| 157 | + content += f":::{{{include}}} {make_relative(file)}\n:::\n" |
| 158 | + content += ":::\n\n" |
| 159 | + |
| 160 | + return content |
| 161 | + |
| 162 | + |
| 163 | +def generate_examples(): |
| 164 | + # Create the EXAMPLE_DOC_DIR if it doesn't exist |
| 165 | + if not EXAMPLE_DOC_DIR.exists(): |
| 166 | + EXAMPLE_DOC_DIR.mkdir(parents=True) |
| 167 | + |
| 168 | + # Create empty indices |
| 169 | + examples_index = Index( |
| 170 | + path=EXAMPLE_DOC_DIR / "examples_index.md", |
| 171 | + title="Examples", |
| 172 | + description= |
| 173 | + "A collection of examples demonstrating usage of vLLM-Spyre.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.", # noqa: E501 |
| 174 | + caption="Examples", |
| 175 | + maxdepth=2) |
| 176 | + # Category indices stored in reverse order because they are inserted into |
| 177 | + # examples_index.documents at index 0 in order |
| 178 | + category_indices = { |
| 179 | + "other": |
| 180 | + Index( |
| 181 | + path=EXAMPLE_DOC_DIR / "examples_other_index.md", |
| 182 | + title="Other", |
| 183 | + description= |
| 184 | + "Other examples that don't strongly fit into the online or offline serving categories.", # noqa: E501 |
| 185 | + caption="Examples", |
| 186 | + ), |
| 187 | + "online_serving": |
| 188 | + Index( |
| 189 | + path=EXAMPLE_DOC_DIR / "examples_online_serving_index.md", |
| 190 | + title="Online Serving", |
| 191 | + description= |
| 192 | + "Online serving examples demonstrate how to use vLLM-Spyre in an online setting, where the model is queried for predictions in real-time.", # noqa: E501 |
| 193 | + caption="Examples", |
| 194 | + ), |
| 195 | + "offline_inference": |
| 196 | + Index( |
| 197 | + path=EXAMPLE_DOC_DIR / "examples_offline_inference_index.md", |
| 198 | + title="Offline Inference", |
| 199 | + description= |
| 200 | + "Offline inference examples demonstrate how to use vLLM-Spyre in an offline setting, where the model is queried for predictions in batches.", # noqa: E501 |
| 201 | + caption="Examples", |
| 202 | + ), |
| 203 | + } |
| 204 | + |
| 205 | + examples = [] |
| 206 | + glob_patterns = ["*.py", "*.md", "*.sh"] |
| 207 | + # Find categorised examples |
| 208 | + for category in category_indices: |
| 209 | + category_dir = EXAMPLE_DIR / category |
| 210 | + globs = [category_dir.glob(pattern) for pattern in glob_patterns] |
| 211 | + for path in itertools.chain(*globs): |
| 212 | + examples.append(Example(path, category)) |
| 213 | + # Find examples in subdirectories |
| 214 | + for path in category_dir.glob("*/*.md"): |
| 215 | + examples.append(Example(path.parent, category)) |
| 216 | + # Find uncategorised examples |
| 217 | + globs = [EXAMPLE_DIR.glob(pattern) for pattern in glob_patterns] |
| 218 | + for path in itertools.chain(*globs): |
| 219 | + examples.append(Example(path)) |
| 220 | + # Find examples in subdirectories |
| 221 | + for path in EXAMPLE_DIR.glob("*/*.md"): |
| 222 | + # Skip categorised examples |
| 223 | + if path.parent.name in category_indices: |
| 224 | + continue |
| 225 | + examples.append(Example(path.parent)) |
| 226 | + |
| 227 | + # Generate the example documentation |
| 228 | + for example in sorted(examples, key=lambda e: e.path.stem): |
| 229 | + doc_path = EXAMPLE_DOC_DIR / f"{example.path.stem}.md" |
| 230 | + with open(doc_path, "w+") as f: |
| 231 | + f.write(example.generate()) |
| 232 | + # Add the example to the appropriate index |
| 233 | + index = category_indices.get(example.category, examples_index) |
| 234 | + index.documents.append(example.path.stem) |
| 235 | + |
| 236 | + # Generate the index files |
| 237 | + for category_index in category_indices.values(): |
| 238 | + if category_index.documents: |
| 239 | + examples_index.documents.insert(0, category_index.path.name) |
| 240 | + with open(category_index.path, "w+") as f: |
| 241 | + f.write(category_index.generate()) |
| 242 | + |
| 243 | + with open(examples_index.path, "w+") as f: |
| 244 | + f.write(examples_index.generate()) |
0 commit comments