+
+ def fill_cache(self):
+ while True:
+ prompts, answers = self.generate_prompts_and_answers_(self.chunk_size)
+
+ self.queue.put((prompts.to("cpu"), answers.to("cpu")), block=True)
+
+ def generate_prompts_and_answers(self, nb):
+ if self.queue is None:
+ return self.generate_prompts_and_answers_(nb)
+
+ if self.rest is not None:
+ prompts, answers = rest
+ else:
+ prompts, answers = [], []
+
+ self.rest = None
+
+ n = sum([p.size(0) for p in prompts])
+
+ with tqdm.tqdm(
+ total=nb,
+ dynamic_ncols=True,
+ desc="world generation",
+ ) as pbar:
+ while n < nb:
+ p, s = self.queue.get(block=True)
+ prompts.append(p)
+ answers.append(s)
+ n += p.size(0)
+ pbar.update(p.size(0))
+
+ prompts, answers = torch.cat(prompts, dim=0), torch.cat(answers, dim=0)
+ assert n == prompts.size(0)
+
+ k = n - nb
+
+ if k > 0:
+ rest = (prompts[-k:], answers[-k:])
+ prompts, answers = prompts[:-k], answers[:-k]
+
+ return prompts, answers