| 60 | |
| 61 | |
| 62 | class LazyWriter: |
| 63 | def __init__(self, path, data_type, is_array=False, array_data_type=np.int32): |
| 64 | lazypath = get_lazy_path(path) |
| 65 | if not os.path.exists(lazypath): |
| 66 | os.makedirs(lazypath) |
| 67 | self.datapath = os.path.join(lazypath, data_type) |
| 68 | self.lenpath = os.path.join(lazypath, data_type + '.len.pkl') |
| 69 | self.array_data_type = array_data_type |
| 70 | self.output = open(self.datapath, 'wb') |
| 71 | self.lengths = [] |
| 72 | self.is_array = is_array |
| 73 | |
| 74 | @staticmethod |
| 75 | def get_len_path(path, data_type): |
| 76 | lazypath = get_lazy_path(path) |
| 77 | return os.path.join(lazypath, data_type + '.len.pkl') |
| 78 | |
| 79 | def write(self, s): |
| 80 | if isinstance(s, dict): |
| 81 | s = s['text'] |
| 82 | if self.is_array: |
| 83 | encoded = np.array(s, dtype=self.array_data_type).tobytes(order='C') |
| 84 | self.output.write(encoded) |
| 85 | self.lengths.append(len(s)) |
| 86 | else: |
| 87 | encoded = s.encode('utf-8') |
| 88 | self.output.write(encoded) |
| 89 | self.lengths.append(len(encoded)) |
| 90 | |
| 91 | def close(self): |
| 92 | self.output.close() |
| 93 | with open(self.lenpath, 'wb') as f: |
| 94 | pkl.dump(self.lengths, f) |
| 95 | |
| 96 | |
| 97 | def split_strings(strings, start, chr_lens): |