| 31 | |
| 32 | |
| 33 | class CodeGeeXTokenizer(object): |
| 34 | def __init__( |
| 35 | self, |
| 36 | tokenizer: GPT2TokenizerFast = None, |
| 37 | tokenizer_path: str = "EleutherAI/gpt-j-6B", |
| 38 | start_extra_id: int = 10, |
| 39 | max_len : int = 10, |
| 40 | mode='codegeex-13b', |
| 41 | dict_file: str = None, |
| 42 | ): |
| 43 | self.tokenizer = tokenizer if tokenizer is not None else AutoTokenizer.from_pretrained(tokenizer_path) |
| 44 | if mode not in ['codegeex-13b']: |
| 45 | raise ValueError(f"Invalid mode {mode}, choose from ['codegeex-13b']") |
| 46 | self.start_extra_id = start_extra_id |
| 47 | self.max_len = max_len |
| 48 | self.mode = mode |
| 49 | self.eos_token_id = self.tokenizer.eos_token_id |
| 50 | |
| 51 | def encode_code(self, code: str): |
| 52 | if self.mode == 'codegeex-13b': |
| 53 | code = encode_whitespaces(code, self.start_extra_id, self.max_len) |
| 54 | input_ids = self.tokenizer(code, is_split_into_words=False, verbose=False).input_ids |
| 55 | |
| 56 | return input_ids |
| 57 | |
| 58 | def decode_code(self, input_ids): |
| 59 | if self.mode == 'codegeex-13b': |
| 60 | text = self.tokenizer.decode(input_ids, skip_special_tokens=False, verbose=False) |
| 61 | output_code = decode_whitespaces(text, self.start_extra_id, self.max_len) |
| 62 | |
| 63 | return output_code |