refactor : pipeline 策略化拆分,消除 _flush if/else
- PackingStrategy / PositionIdStrategy / StoreWriter 独立文件 + Factory - Pipeline._flush 零 if/else,纯编排 - SectionRenderer 从 SectionedMaskBuilder 分离 - OutputConfig.position_ids_mode 默认改为 ""none""
This commit is contained in:
@@ -0,0 +1,50 @@
|
||||
"""Position-id generation strategies for packed sequences.
|
||||
|
||||
Each strategy takes the list of per-document token sequences after packing
|
||||
and returns a flat list of position ids (same total length as all
|
||||
sequences combined). The pipeline wraps the result into a tensor and
|
||||
attaches it as ``position_ids``.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List
|
||||
|
||||
from astrai.factory import BaseFactory
|
||||
|
||||
|
||||
class PositionIdStrategy(ABC):
|
||||
"""Generate ``position_ids`` for packed sequences."""
|
||||
|
||||
@abstractmethod
|
||||
def generate(self, sequences: List[list]) -> List[int]: ...
|
||||
|
||||
|
||||
class PositionIdStrategyFactory(BaseFactory["PositionIdStrategy"]):
|
||||
@classmethod
|
||||
def _validate_component(cls, component_cls: type):
|
||||
if not issubclass(component_cls, PositionIdStrategy):
|
||||
raise TypeError(
|
||||
f"{component_cls.__name__} must inherit from PositionIdStrategy"
|
||||
)
|
||||
|
||||
|
||||
@PositionIdStrategyFactory.register("none")
|
||||
class NoPositionId(PositionIdStrategy):
|
||||
def generate(self, sequences):
|
||||
return []
|
||||
|
||||
|
||||
@PositionIdStrategyFactory.register("doc_reset")
|
||||
class DocResetPositionId(PositionIdStrategy):
|
||||
def generate(self, sequences):
|
||||
pos_ids = []
|
||||
for seq in sequences:
|
||||
pos_ids.extend(range(len(seq)))
|
||||
return pos_ids
|
||||
|
||||
|
||||
@PositionIdStrategyFactory.register("continuous")
|
||||
class ContinuousPositionId(PositionIdStrategy):
|
||||
def generate(self, sequences):
|
||||
total = sum(len(seq) for seq in sequences)
|
||||
return list(range(total))
|
||||
Reference in New Issue
Block a user