| 1026 | |
| 1027 | |
| 1028 | class SequenceHistory: |
| 1029 | __slots__ = ("_position_lengths", "_root", "_tails", "size") |
| 1030 | |
| 1031 | @dataclass |
| 1032 | class Node: |
| 1033 | token: Optional[int] = None |
| 1034 | parent: Optional["SequenceHistory.Node"] = None |
| 1035 | children: Dict[int, "SequenceHistory.Node"] = field(default_factory=dict) |
| 1036 | sequences: set[int] = field(default_factory=set) |
| 1037 | position_increment: int = 1 |
| 1038 | |
| 1039 | def __init__(self) -> None: |
| 1040 | self._root = SequenceHistory.Node() |
| 1041 | self._tails: Dict[int, SequenceHistory.Node] = {} |
| 1042 | self._position_lengths: Dict[int, int] = {} |
| 1043 | self.size = 0 |
| 1044 | |
| 1045 | def extend( |
| 1046 | self, |
| 1047 | sequence_id: int, |
| 1048 | tokens: Sequence[int], |
| 1049 | position_increments: Optional[Sequence[int]] = None, |
| 1050 | ) -> None: |
| 1051 | assert sequence_id >= 0 |
| 1052 | if position_increments is None: |
| 1053 | position_increments = [1] * len(tokens) |
| 1054 | assert len(position_increments) == len(tokens) |
| 1055 | node = self._tails.get(sequence_id, self._root) |
| 1056 | position_length = self._position_lengths.get(sequence_id, 0) |
| 1057 | for token, position_increment in zip(tokens, position_increments): |
| 1058 | position_increment = max(0, int(position_increment)) |
| 1059 | child = node.children.get(sequence_id) |
| 1060 | if child is None: |
| 1061 | child = SequenceHistory.Node( |
| 1062 | token=token, |
| 1063 | parent=node, |
| 1064 | position_increment=position_increment, |
| 1065 | ) |
| 1066 | node.children[sequence_id] = child |
| 1067 | self.size += 1 |
| 1068 | else: |
| 1069 | assert child.parent is node |
| 1070 | assert child.token == token |
| 1071 | assert child.position_increment == position_increment |
| 1072 | child.sequences.add(sequence_id) |
| 1073 | position_length += position_increment |
| 1074 | node = child |
| 1075 | if node is self._root: |
| 1076 | self._tails.pop(sequence_id, None) |
| 1077 | self._position_lengths.pop(sequence_id, None) |
| 1078 | else: |
| 1079 | self._tails[sequence_id] = node |
| 1080 | self._position_lengths[sequence_id] = position_length |
| 1081 | |
| 1082 | def position_length(self, sequence_id: int) -> int: |
| 1083 | return self._position_lengths.get(sequence_id, 0) |
| 1084 | |
| 1085 | def position_length_for_prefix(self, sequence_id: int, keep_len: int) -> int: |
no outgoing calls
no test coverage detected
searching dependent graphs…