HF Trainer 사용 시 collator에 아무것도 들어오지 않는 경우

Huggingface Trainer 사용시, dataset과 collator class을 같이 넘겨준다. 

그런데 getitem까지는 잘 데이터가 담기는데, collator의 call 부분을 찍을 시, 아무것도 

들어오지 않는 ( [{}, {}, {}] -> bs 3인 경우) 경우가 생겼을 때

 

-> Trainer에 넘겨주는 training argument 중 remove_unused_columns을 False로 주면 된다. (기본은 True)

 

Trainer 소스 코드에서 아래 함수로 인해 문제가 발생하는 것인데, 

    def _get_collator_with_removed_columns(
        self, data_collator: Callable, description: Optional[str] = None
    ) -> Callable:
        """Wrap the data collator in a callable removing unused columns."""
        if not self.args.remove_unused_columns:
            return data_collator
        self._set_signature_columns_if_needed()
        signature_columns = self._signature_columns

        remove_columns_collator = RemoveColumnsCollator(
            data_collator=data_collator,
            signature_columns=signature_columns,
            logger=logger,
            description=description,
            model_name=self.model.__class__.__name__,
        )
        return remove_columns_collator

 

self.signature_columns은 ['return_loss', 'sample', 'label', 'label_ids'] 이다. 즉, 미리 정해진 key값들이 아닌 다른 key값들은 무시하게 하는데, 내 dataset getitem은 저 중 아무것도 없기에 실제로 collator에 넘기는 것이 없게 된다.