Files
OpenLLM/tools/update-models-import.py
2023-09-01 17:00:49 +00:00

36 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__, globals()["__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',
'from openllm_core.utils import LazyModule',
create_module_import(),
*create_stubs_import(),
]))
return 0
if __name__ == '__main__': raise SystemExit(main())