ViperEkura
fec376b0dd
- 修复 strategy.py 单元素 Union 与缺失的参数/返回类型注解 - 修复 train_context.py 8 个 default=None 字段缺 Optional 标记 - 修复 sample.py/packing.py/position_id.py 方法缺参数及返回类型注解 - 修复 factory.py _resolve_type/list_registered 缺类型注解 - 修复 train_config.py 裸 dict/list 缺泛型参数 - abstractmethod body 从 ... 改为 raise NotImplementedError - feat : checkpoint meta.json 保存 TrainConfig 超参供人工查阅
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
"""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[int]]) -> List[int]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class PositionIdStrategyFactory(BaseFactory["PositionIdStrategy"]):
|
|
pass
|
|
|
|
|
|
@PositionIdStrategyFactory.register("none")
|
|
class NoPositionId(PositionIdStrategy):
|
|
def generate(self, sequences: List[List[int]]) -> List[int]:
|
|
return []
|
|
|
|
|
|
@PositionIdStrategyFactory.register("doc_reset")
|
|
class DocResetPositionId(PositionIdStrategy):
|
|
def generate(self, sequences: List[List[int]]) -> List[int]:
|
|
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: List[List[int]]) -> List[int]:
|
|
total = sum(len(seq) for seq in sequences)
|
|
return list(range(total))
|