File size: 690 Bytes
74b17e0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
import os
from typing import Dict
from .base import *
from ...utils import import_modules
TEMPlATE_FACTORY: Dict[str, Template] = {}
def TemplateFactory(version):
template = TEMPlATE_FACTORY.get(version, None)
assert template, f"{version} is not implmentation"
return template
def register_template(name):
def register_template_cls(cls):
if name in TEMPlATE_FACTORY:
return TEMPlATE_FACTORY[name]
TEMPlATE_FACTORY[name] = cls
return cls
return register_template_cls
# automatically import any Python files in the models/ directory
models_dir = os.path.dirname(__file__)
import_modules(models_dir, "tinyllava.data.template")
|