| | |
| | |
| | |
| | |
| | |
| | |
| |
|
| |
|
| | 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) |
