| |
| |
| |
| |
| |
| |
|
|
|
|
| from typing import Dict, List |
|
|
|
|
| class Trie(object): |
| def __init__(self, sequences: List[List[int]] = []): |
| self.trie_dict = {} |
| self.len = 0 |
| if sequences: |
| for sequence in sequences: |
| Trie._add_to_trie(sequence, self.trie_dict) |
| self.len += 1 |
|
|
| self.append_trie = None |
| self.bos_token_id = None |
|
|
| def append(self, trie, bos_token_id): |
| self.append_trie = trie |
| self.bos_token_id = bos_token_id |
|
|
| def add(self, sequence: List[int]): |
| Trie._add_to_trie(sequence, self.trie_dict) |
| self.len += 1 |
|
|
| def get(self, prefix_sequence: List[int]): |
| return Trie._get_from_trie( |
| prefix_sequence, self.trie_dict, self.append_trie, self.bos_token_id |
| ) |
|
|
| @staticmethod |
| def load_from_dict(trie_dict): |
| trie = Trie() |
| trie.trie_dict = trie_dict |
| trie.len = sum(1 for _ in trie) |
| return trie |
|
|
| @staticmethod |
| def _add_to_trie(sequence: List[int], trie_dict: Dict): |
| if sequence: |
| if sequence[0] not in trie_dict: |
| trie_dict[sequence[0]] = {} |
| Trie._add_to_trie(sequence[1:], trie_dict[sequence[0]]) |
|
|
| @staticmethod |
| def _get_from_trie( |
| prefix_sequence: List[int], |
| trie_dict: Dict, |
| append_trie=None, |
| bos_token_id: int = None, |
| ): |
| if len(prefix_sequence) == 0: |
| output = list(trie_dict.keys()) |
| if append_trie and bos_token_id in output: |
| output.remove(bos_token_id) |
| output += list(append_trie.trie_dict.keys()) |
| return output |
| elif prefix_sequence[0] in trie_dict: |
| return Trie._get_from_trie( |
| prefix_sequence[1:], |
| trie_dict[prefix_sequence[0]], |
| append_trie, |
| bos_token_id, |
| ) |
| else: |
| if append_trie: |
| return append_trie.get(prefix_sequence) |
| else: |
| return [] |
|
|
| def __iter__(self): |
| def _traverse(prefix_sequence, trie_dict): |
| if trie_dict: |
| for next_token in trie_dict: |
| yield from _traverse( |
| prefix_sequence + [next_token], trie_dict[next_token] |
| ) |
| else: |
| yield prefix_sequence |
|
|
| return _traverse([], self.trie_dict) |
|
|
| def __len__(self): |
| return self.len |
|
|
| def __getitem__(self, value): |
| return self.get(value) |
