| 31 | os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))) |
| 32 | |
| 33 | class Encoder(object): |
| 34 | def __init__(self, args): |
| 35 | self.args = args |
| 36 | |
| 37 | def initializer(self): |
| 38 | # Use Encoder class as a container for global data |
| 39 | Encoder.tokenizer = build_tokenizer(self.args) |
| 40 | |
| 41 | def encode(self, text): |
| 42 | if self.args.ftfy: |
| 43 | text = ftfy.fix_text(text) |
| 44 | ids = {} |
| 45 | for key in self.args.jsonl_keys: |
| 46 | doc_ids = [] |
| 47 | try: |
| 48 | text_ids = Encoder.tokenizer(text, add_special_tokens=False)['input_ids'] |
| 49 | """ |
| 50 | text_ids = Encoder.tokenizer(text, add_special_tokens=False, padding='max_length', |
| 51 | max_length=2047, truncation=True)['input_ids'] |
| 52 | """ |
| 53 | if max(text_ids) >= Encoder.tokenizer.vocab_size: |
| 54 | print(text) |
| 55 | print(max(text_ids)) |
| 56 | continue |
| 57 | except: |
| 58 | continue |
| 59 | if len(text_ids) > 0: |
| 60 | doc_ids.append(text_ids) |
| 61 | if self.args.append_eod: |
| 62 | if hasattr(Encoder.tokenizer, 'eos_token_id'): |
| 63 | doc_ids[-1].append(Encoder.tokenizer.eos_token_id) |
| 64 | elif hasattr(Encoder.tokenizer, 'eod_id'): |
| 65 | doc_ids[-1].append(Encoder.tokenizer.eod_id) |
| 66 | else: |
| 67 | doc_ids[-1].append(Encoder.tokenizer.eod) |
| 68 | #doc_ids[-1].append(Encoder.tokenizer.pad_token_id) |
| 69 | ids[key] = doc_ids |
| 70 | return ids, len(text) |
| 71 | |
| 72 | def get_args(): |
| 73 | parser = argparse.ArgumentParser() |