| |
|
|
| from typing import Dict, Literal, Optional |
|
|
| from ..utils import Processor |
| from .base import Template |
| from .template_meta import TemplateMeta |
|
|
| TEMPLATE_MAPPING: Dict[str, TemplateMeta] = {} |
|
|
|
|
| def register_template(template_meta: TemplateMeta, *, exist_ok: bool = False) -> None: |
| template_type = template_meta.template_type |
| if not exist_ok and template_type in TEMPLATE_MAPPING: |
| raise ValueError(f'The `{template_type}` has already been registered in the TEMPLATE_MAPPING.') |
| TEMPLATE_MAPPING[template_type] = template_meta |
|
|
|
|
| def get_template( |
| template_type: str, |
| processor: Processor, |
| default_system: Optional[str] = None, |
| max_length: Optional[int] = None, |
| *, |
| use_chat_template: bool = True, |
| template_backend: Literal['swift', 'jinja'] = 'swift', |
| truncation_strategy: Literal['raise', 'left', 'right'] = 'raise', |
| max_pixels: Optional[int] = None, |
| agent_template: Optional[str] = None, |
| norm_bbox: Literal['norm1000', 'none', None] = None, |
| response_prefix: Optional[str] = None, |
| |
| padding_side: Literal['left', 'right'] = 'right', |
| loss_scale: str = 'default', |
| sequence_parallel_size: int = 1) -> 'Template': |
| template_meta = TEMPLATE_MAPPING[template_type] |
| template_cls = template_meta.template_cls |
| return template_cls( |
| processor, |
| template_meta, |
| default_system, |
| max_length, |
| use_chat_template=use_chat_template, |
| template_backend=template_backend, |
| truncation_strategy=truncation_strategy, |
| max_pixels=max_pixels, |
| agent_template=agent_template, |
| norm_bbox=norm_bbox, |
| response_prefix=response_prefix, |
| padding_side=padding_side, |
| loss_scale=loss_scale, |
| sequence_parallel_size=sequence_parallel_size, |
| ) |
|
|
|
|
| def get_template_meta(template_type: str) -> TemplateMeta: |
| return TEMPLATE_MAPPING[template_type] |
|
|
