123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
"""Factory for creating and registering processors."""
|
|
from typing import Dict, List, Any, Optional, Type
|
|
|
|
from pipeline.processors.base import BaseProcessor
|
|
from pipeline.tokenizer import BpeTokenizer
|
|
from pipeline.strategies import PromptStrategy, StrategyFactory
|
|
|
|
|
|
class ProcessorFactory:
|
|
"""Registry and factory for BaseProcessor implementations.
|
|
|
|
Supports decorator-based registration for extensible processor types.
|
|
|
|
Example usage::
|
|
|
|
@ProcessorFactory.register("custom")
|
|
class CustomProcessor(BaseProcessor):
|
|
...
|
|
|
|
processor = ProcessorFactory.create(optimizer, "custom", **kwargs)
|
|
"""
|
|
|
|
PROCESSOR_MAP: Dict[str, Type[BaseProcessor]] = {}
|
|
|
|
@classmethod
|
|
def register(cls, name: str):
|
|
"""Decorator to register a new processor class.
|
|
|
|
Args:
|
|
name: Registration name for the processor.
|
|
|
|
Returns:
|
|
Decorator function that registers the processor class.
|
|
"""
|
|
|
|
def decorator(processor_cls: Type[BaseProcessor]) -> Type[BaseProcessor]:
|
|
if not issubclass(processor_cls, BaseProcessor):
|
|
raise TypeError(
|
|
f"{processor_cls.__name__} must inherit from BaseProcessor"
|
|
)
|
|
cls.PROCESSOR_MAP[name] = processor_cls
|
|
return processor_cls
|
|
|
|
return decorator
|
|
|
|
@classmethod
|
|
def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor:
|
|
"""Create a processor by type name (uses default ChatMLStrategy for SFT/DPO).
|
|
|
|
Args:
|
|
processor_type: Registered processor name (e.g. ``"pt"``, ``"sft"``, ``"dpo"``).
|
|
tokenizer: Tokenizer instance.
|
|
|
|
Returns:
|
|
Processor instance.
|
|
|
|
Raises:
|
|
ValueError: If processor_type is not registered.
|
|
"""
|
|
if processor_type not in cls.PROCESSOR_MAP:
|
|
raise ValueError(
|
|
f"Unknown processor type: '{processor_type}'. "
|
|
f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
|
|
)
|
|
return cls.PROCESSOR_MAP[processor_type](tokenizer)
|
|
|
|
@classmethod
|
|
def create_with_strategy(
|
|
cls,
|
|
processor_type: str,
|
|
tokenizer: BpeTokenizer,
|
|
strategy: PromptStrategy,
|
|
) -> BaseProcessor:
|
|
"""Create a processor with a custom strategy.
|
|
|
|
Only SFT and DPO processors accept a strategy; PreTrain ignores it.
|
|
|
|
Args:
|
|
processor_type: Registered processor name.
|
|
tokenizer: Tokenizer instance.
|
|
strategy: Prompt strategy instance.
|
|
|
|
Returns:
|
|
Processor instance configured with strategy.
|
|
"""
|
|
if processor_type not in cls.PROCESSOR_MAP:
|
|
raise ValueError(
|
|
f"Unknown processor type: '{processor_type}'. "
|
|
f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
|
|
)
|
|
|
|
processor_cls = cls.PROCESSOR_MAP[processor_type]
|
|
if processor_type == "pt":
|
|
return processor_cls(tokenizer)
|
|
return processor_cls(tokenizer, strategy=strategy)
|
|
|
|
@classmethod
|
|
def create_with_strategy_name(
|
|
cls,
|
|
processor_type: str,
|
|
tokenizer: BpeTokenizer,
|
|
strategy_name: str,
|
|
**strategy_kwargs,
|
|
) -> BaseProcessor:
|
|
"""Create a processor with a strategy selected by name.
|
|
|
|
Args:
|
|
processor_type: Registered processor name.
|
|
tokenizer: Tokenizer instance.
|
|
strategy_name: Registered strategy name (``"chatml"``, ``"alpaca"``, etc.).
|
|
**strategy_kwargs: Forwarded to the strategy constructor.
|
|
|
|
Returns:
|
|
Processor instance.
|
|
"""
|
|
strategy = StrategyFactory.create(strategy_name, tokenizer, **strategy_kwargs)
|
|
return cls.create_with_strategy(processor_type, tokenizer, strategy)
|
|
|
|
@classmethod
|
|
def available_types(cls) -> List[str]:
|
|
"""Return list of registered processor type names."""
|
|
return list(cls.PROCESSOR_MAP.keys())
|