Remove specified columns from the input mini-batch, and convert them to attributes. For example: if columns_to_remove = ['id'], then user should call batch.id instead of batch['id']. Args: data_collator: An inner data collator to collate the mini-batch columns_to_remove(`Li
| 11 | |
| 12 | |
| 13 | class RemoveColumnsCollator: |
| 14 | """Remove specified columns from the input mini-batch, and convert them to attributes. |
| 15 | |
| 16 | For example: if columns_to_remove = ['id'], then user should call batch.id instead of batch['id']. |
| 17 | |
| 18 | Args: |
| 19 | data_collator: An inner data collator to collate the mini-batch |
| 20 | columns_to_remove(`List[str]`): The redundant columns to be removed from the mini-batch |
| 21 | model_name(`Optional[str]`): An optional model name to print into log |
| 22 | description(`Optional[str]`): An optional description to print into log |
| 23 | """ |
| 24 | |
| 25 | def __init__( |
| 26 | self, |
| 27 | data_collator, |
| 28 | columns_to_remove: List[str], |
| 29 | model_name: Optional[str] = None, |
| 30 | description: Optional[str] = None, |
| 31 | ): |
| 32 | self.data_collator = data_collator |
| 33 | self.columns_to_remove = columns_to_remove |
| 34 | self.description = description |
| 35 | self.model_name = model_name |
| 36 | self.message_logged = False |
| 37 | |
| 38 | def _remove_columns(self, feature: Mapping) -> Tuple[Mapping, Any]: |
| 39 | if not isinstance(feature, Mapping): |
| 40 | return feature, None |
| 41 | if not self.message_logged and self.model_name: |
| 42 | ignored_columns = list( |
| 43 | set(feature.keys()) - set(self.columns_to_remove)) |
| 44 | if len(ignored_columns) > 0: |
| 45 | dset_description = '' if self.description is None else f'in the {self.description} set' |
| 46 | logger.info( |
| 47 | f"The following columns {dset_description} don't have a corresponding argument in " |
| 48 | f"`{self.model_name}.forward` and have been ignored: {', '.join(ignored_columns)}." |
| 49 | f"Legal columns: {', '.join(self.columns_to_remove)}." |
| 50 | f" If {', '.join(ignored_columns)} are not expected by `{self.model_name}.forward`, " |
| 51 | ' you can safely ignore this message.') |
| 52 | self.message_logged = True |
| 53 | feature_clean = { |
| 54 | k: v |
| 55 | for k, v in feature.items() if k in self.columns_to_remove |
| 56 | } |
| 57 | feature_unused = { |
| 58 | k: v |
| 59 | for k, v in feature.items() if k not in self.columns_to_remove |
| 60 | } |
| 61 | return feature_clean, feature_unused |
| 62 | |
| 63 | def __call__(self, features: List[Mapping]): |
| 64 | features_clean = [] |
| 65 | features_unused = [] |
| 66 | for feature in features: |
| 67 | feature, feature_unused = self._remove_columns(feature) |
| 68 | features_clean.append(feature) |
| 69 | features_unused.append(feature_unused) |
| 70 | data = OrderedDict(self.data_collator(features_clean)) |
no outgoing calls
no test coverage detected
searching dependent graphs…