Files
OpenLLM/tools/update-models-import.py
2023-08-15 02:11:14 -04:00

27 lines
1.4 KiB
Python
Executable File

#!/usr/bin/env python3
from __future__ import annotations
import os
from pathlib import Path
_TARGET_FILE = Path(__file__).parent.parent/"openllm-python"/"src"/"openllm"/"models"/"__init__.py"
def create_module_import() -> str:
r = [f'"{p.name}"' for p in _TARGET_FILE.parent.glob('*/') if p.name not in ['__pycache__', '__init__.py', '.DS_Store']]
return f"_MODELS: set[str] = {{{', '.join(sorted(r))}}}"
def create_stubs_import() -> list[str]: return ["if t.TYPE_CHECKING: from . import "+",".join([f"{p.name} as {p.name}" for p in sorted(_TARGET_FILE.parent.glob("*/")) if p.name not in {"__pycache__", "__init__.py", ".DS_Store"}]),
'__lazy=LazyModule(__name__, os.path.abspath("__file__"), {k: [] for k in _MODELS})', "__all__=__lazy.__all__", "__dir__=__lazy.__dir__", "__getattr__=__lazy.__getattr__\n"]
def main() -> int:
_path = os.path.join(os.path.basename(os.path.dirname(__file__)), os.path.basename(__file__))
with _TARGET_FILE.open("w") as f: f.writelines("\n".join([
f"# This file is generated by {_path}. DO NOT EDIT MANUALLY!",
f"# To update this, run ./{_path}",
"from __future__ import annotations",
"import typing as t, os",
"from openllm.utils import LazyModule",
create_module_import(),
*create_stubs_import(),
]))
return 0
if __name__ == "__main__": raise SystemExit(main())