(
self,
sequence_id: int,
tokens: Sequence[int],
)
| 876 | return self.root |
| 877 | |
| 878 | def extend( |
| 879 | self, |
| 880 | sequence_id: int, |
| 881 | tokens: Sequence[int], |
| 882 | ) -> None: |
| 883 | assert sequence_id >= 0 |
| 884 | node = self.sequences.get(sequence_id, self.root) |
| 885 | if tokens: |
| 886 | node.tails.discard(sequence_id) |
| 887 | length = self.sequence_lengths.get(sequence_id, 0) |
| 888 | index = 0 |
| 889 | while index < len(tokens): |
| 890 | token = tokens[index] |
| 891 | child = node.children.get(token) |
| 892 | if child is None: |
| 893 | child = RadixTrie.Node( |
| 894 | label=tuple(tokens[index:]), |
| 895 | parent=node, |
| 896 | sequences={sequence_id}, |
| 897 | ) |
| 898 | node.children[token] = child |
| 899 | node = child |
| 900 | length += len(tokens) - index |
| 901 | index = len(tokens) |
| 902 | break |
| 903 | match_len = self._common_prefix_len(child.label, tokens, index) |
| 904 | if match_len == len(child.label): |
| 905 | child.sequences.add(sequence_id) |
| 906 | node = child |
| 907 | length += match_len |
| 908 | index += match_len |
| 909 | continue |
| 910 | node = self._split_child(node, child, match_len) |
| 911 | node.sequences.add(sequence_id) |
| 912 | length += match_len |
| 913 | index += match_len |
| 914 | if index == len(tokens): |
| 915 | break |
| 916 | suffix = RadixTrie.Node( |
| 917 | label=tuple(tokens[index:]), |
| 918 | parent=node, |
| 919 | sequences={sequence_id}, |
| 920 | ) |
| 921 | node.children[suffix.label[0]] = suffix |
| 922 | node = suffix |
| 923 | length += len(tokens) - index |
| 924 | index = len(tokens) |
| 925 | break |
| 926 | if node is self.root: |
| 927 | self.sequences.pop(sequence_id, None) |
| 928 | self.sequence_lengths.pop(sequence_id, None) |
| 929 | self.root.tails.discard(sequence_id) |
| 930 | else: |
| 931 | self.sequences[sequence_id] = node |
| 932 | self.sequence_lengths[sequence_id] = length |
| 933 | node.tails.add(sequence_id) |
| 934 | |
| 935 | def length(self, sequence_id: int) -> int: |
no test coverage detected