| 74 | |
| 75 | |
| 76 | class ConstraintDecoder: |
| 77 | def __init__(self, tokenizer, source_prefix): |
| 78 | self.tokenizer = tokenizer |
| 79 | self.source_prefix = source_prefix |
| 80 | self.source_prefix_tokenized = tokenizer.encode(source_prefix, |
| 81 | add_special_tokens=False) if source_prefix else [] |
| 82 | |
| 83 | def get_state_valid_tokens(self, src_sentence: List[str], tgt_generated: List[str]) -> List[str]: |
| 84 | pass |
| 85 | |
| 86 | def constraint_decoding(self, src_sentence, tgt_generated): |
| 87 | if self.source_prefix_tokenized: |
| 88 | # Remove Source Prefix for Generation |
| 89 | src_sentence = src_sentence[len(self.source_prefix_tokenized):] |
| 90 | |
| 91 | if debug: |
| 92 | print("Src:", self.tokenizer.convert_ids_to_tokens(src_sentence)) |
| 93 | print("Tgt:", self.tokenizer.convert_ids_to_tokens(tgt_generated)) |
| 94 | |
| 95 | valid_token_ids = self.get_state_valid_tokens( |
| 96 | src_sentence.tolist(), |
| 97 | tgt_generated.tolist() |
| 98 | ) |
| 99 | |
| 100 | if debug: |
| 101 | print('========================================') |
| 102 | print('valid tokens:', self.tokenizer.convert_ids_to_tokens( |
| 103 | valid_token_ids), valid_token_ids) |
| 104 | if debug_step: |
| 105 | input() |
| 106 | |
| 107 | # return self.tokenizer.convert_tokens_to_ids(valid_tokens) |
| 108 | return valid_token_ids |
| 109 | |
| 110 | # ET + RT + Src -> ((Role)(Role)), ETRTText2Role 使用 |
| 111 | class RoleConstraintDecoder(ConstraintDecoder): |
nothing calls this directly
no outgoing calls
no test coverage detected