MCPcopy Create free account
hub / github.com/THUDM/GLM / LazyLoader

Class LazyLoader

data_utils/lazy_loader.py:122–262  ·  view source on GitHub ↗

Arguments: path: path to directory where array entries are concatenated into one big string file and the .len file are located data_type (str): Some datsets have multiple fields that are stored in different paths. `data_type` specifies which of these fiel

Source from the content-addressed store, hash-verified

120
121
122class LazyLoader(object):
123 """
124 Arguments:
125 path: path to directory where array entries are concatenated into one big string file
126 and the .len file are located
127 data_type (str): Some datsets have multiple fields that are stored in different paths.
128 `data_type` specifies which of these fields to load in this class
129 mem_map (boolean): Specifies whether to memory map file `path`
130 map_fn (callable): Fetched strings are passed through map_fn before being returned.
131
132 Example of lazy loader directory structure:
133 file.json
134 file.lazy/
135 data_type1
136 data_type1.len.pkl
137 data_type2
138 data_type2.len.pkl
139 """
140
141 def __init__(self, path, data_type='data', mem_map=False, map_fn=None, is_array=False, array_data_type=np.int32,
142 load_memory=False, half_load=False):
143 lazypath = get_lazy_path(path)
144 datapath = os.path.join(lazypath, data_type)
145 # get file where array entries are concatenated into one big string
146 self._file = open(datapath, 'rb')
147 self.file = self._file
148 self.is_array = is_array
149 self.array_data_type = array_data_type
150 # memory map file if necessary
151 lenpath = os.path.join(lazypath, data_type + '.len.pkl')
152 self.lens = pkl.load(open(lenpath, 'rb'))
153 if half_load:
154 self.lens = self.lens[:2 * len(self.lens) // 3]
155 self.ends = list(accumulate(self.lens))
156 self.dumb_ends = list(self.ends)
157 self.mem_map = mem_map
158 self.load_memory = load_memory
159 if self.load_memory:
160 data_type_size = np.dtype(self.array_data_type).itemsize
161 if half_load:
162 self.file = self.file.read(sum(self.lens) * data_type_size)
163 else:
164 self.file = self.file.read()
165 self.file = np.ndarray(shape=(len(self.file) // data_type_size,), dtype=array_data_type, buffer=self.file,
166 order='C')
167 elif self.mem_map:
168 if is_array:
169 if self.ends[-1] == 0:
170 self.file = np.array([], dtype=array_data_type)
171 else:
172 self.file = np.memmap(self.file, dtype=array_data_type, mode='r', order='C')
173 else:
174 if self.ends[-1] == 0:
175 self.file = bytearray()
176 else:
177 self.file = mmap.mmap(self.file.fileno(), 0, prot=mmap.PROT_READ)
178 self.read_lock = Lock()
179 self.process_fn = map_fn

Callers 1

get_datasetFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected