Select a subset of the DataProto via batch_keys and meta_info_keys Args: batch_keys (list, optional): a list of strings indicating the keys in batch to select meta_info_keys (list, optional): a list of keys indicating the meta info to select Returns:
(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False)
| 609 | return self |
| 610 | |
| 611 | def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> "DataProto": |
| 612 | """Select a subset of the DataProto via batch_keys and meta_info_keys |
| 613 | |
| 614 | Args: |
| 615 | batch_keys (list, optional): a list of strings indicating the keys in batch to select |
| 616 | meta_info_keys (list, optional): a list of keys indicating the meta info to select |
| 617 | |
| 618 | Returns: |
| 619 | DataProto: the DataProto with the selected batch_keys and meta_info_keys |
| 620 | """ |
| 621 | # TODO (zhangchi.usc1992) whether to copy |
| 622 | if batch_keys is not None: |
| 623 | batch_keys = tuple(batch_keys) |
| 624 | sub_batch = self.batch.select(*batch_keys) |
| 625 | else: |
| 626 | sub_batch = self.batch |
| 627 | |
| 628 | if non_tensor_batch_keys is not None: |
| 629 | non_tensor_batch = {key: val for key, val in self.non_tensor_batch.items() if key in non_tensor_batch_keys} |
| 630 | else: |
| 631 | non_tensor_batch = self.non_tensor_batch |
| 632 | |
| 633 | if deepcopy: |
| 634 | non_tensor_batch = copy.deepcopy(non_tensor_batch) |
| 635 | |
| 636 | if meta_info_keys is not None: |
| 637 | sub_meta_info = {key: val for key, val in self.meta_info.items() if key in meta_info_keys} |
| 638 | else: |
| 639 | sub_meta_info = self.meta_info |
| 640 | |
| 641 | if deepcopy: |
| 642 | sub_meta_info = copy.deepcopy(sub_meta_info) |
| 643 | |
| 644 | return type(self)(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info) |
| 645 | |
| 646 | def select_idxs(self, idxs): |
| 647 | """ |
no outgoing calls