Method items

lmdeploy/turbomind/loader.py:31–44

(self)

        self.pattern = pattern

    def items(self):
        for data in iter(self.que.get, None):
            match = []
            if self.pattern:
                for k in data.keys():
                    match = re.findall(self.pattern, k)
                    break
            if not match:
                yield (-1, data)
            else:
                idx = int(match[0])
                yield (idx, data)
            torch.cuda.empty_cache()
            self.que.task_done()

    def all_items(self) -> dict:
        raise NotImplementedError(