| 85 | return json.dumps(output), len(json_line) |
| 86 | |
| 87 | def encode(self, json_line): |
| 88 | data = json.loads(json_line) |
| 89 | ids = {} |
| 90 | lens = {} |
| 91 | for key in self.args.json_keys: |
| 92 | text = data[key] |
| 93 | if isinstance(text, list): |
| 94 | sentences = text |
| 95 | else: |
| 96 | sentences = [text] |
| 97 | doc_ids = [] |
| 98 | sentence_lens = [] |
| 99 | for sentence in sentences: |
| 100 | sentence_ids = Encoder.tokenizer.tokenize(sentence) |
| 101 | if len(sentence_ids) > 0: |
| 102 | doc_ids.extend(sentence_ids) |
| 103 | sentence_lens.append(len(sentence_ids)) |
| 104 | if len(doc_ids) > 0 and self.args.append_eod: |
| 105 | doc_ids.append(Encoder.tokenizer.eod) |
| 106 | sentence_lens[-1] += 1 |
| 107 | ids[key] = doc_ids |
| 108 | lens[key] = sentence_lens |
| 109 | return ids, lens, len(json_line) |
| 110 | |
| 111 | |
| 112 | class Partition(object): |