Week 13 - pickle error

개요

이번 프로젝트를 위해 베이스라인 코드를 모듈화하던 중 발생한 에러이다.

pickle?

텍스트가 아닌 파이썬의 객체를 저장하는 방식이다. 텍스트 데이터를 tokenize 할 때 멀티 프로세싱을 수행하면, 병렬화 해야 하는 대상들을 pickle로 만들어서 세션끼리 공유한다.
그렇게 나도 Hugging Face의 Dataset 함수 중 map을 사용해서 멀티 프로세싱을 하려고 했지만 에러가 발생했다.

기존 코드

class MyModel:
    def __init__(self, config):
        ...
        self.tokenizer = AutoTokenizer.from_pretrained(config["model_name_or_path"])
        ...
    
    def tokenize(self, processed_train):
        # type(processed_train)은 Hugging Face의 Dataset
        def tokenize_fn(element):
            output_texts = []
            for i in range(len(element["messages"])):
                output_texts.append(
                    self.tokenizer.apply_chat_template(
                        element["messages"][i],
                        tokenize=False,
                    )
                )

            outputs = self.tokenizer(
                output_texts,
                truncation=False,
                padding=False,
                return_overflowing_tokens=False,
                return_length=False
            )

            return {
                "input_ids": outputs["input_ids"],
                "attention_mask": outputs["attention_mask"],
            }

        tokenized = processed_train.map(
            tokenize_fn,
            remove_columns=list(processed_train.features),
            batched=True,
            num_proc=4,
            load_from_cache_file=True,
            desc="Tokenizing"
        )
        ...

이렇게 self.tokenizer를 __init__()에서 초기화한 후 tokenize_fn를 병렬화하려고 하면 아래와 같은 에러가 뜬다.

_pickle.PicklingError: Can't pickle <class 'builtins.safe_open'>: it's not found as builtins.safe_open

pickle API 문서 중 What can be pickled and unpickled?에 들어가면 pickle화 할 수 있는 종류들이 나와있다.
내 경우에는 tokenize_fn을 병렬화 할 때 self.tokenizer는 독립적이지 못해서(class에 종속되어) 함께 병렬화되지 못한 것이다.

수정한 코드

def tokenize(self, processed_train):
    tokenizer = self.tokenizer
    # type(processed_train)은 Hugging Face의 Dataset
    def tokenize_fn(element):
        output_texts = []
        for i in range(len(element["messages"])):
            output_texts.append(
                tokenizer.apply_chat_template(
                    element["messages"][i],
                    tokenize=False,
                )
            )

        outputs = tokenizer(
            output_texts,
            truncation=False,
            padding=False,
            return_overflowing_tokens=False,
            return_length=False
        )

        return {
            "input_ids": outputs["input_ids"],
            "attention_mask": outputs["attention_mask"],
        }

    tokenized = processed_train.map(
        tokenize_fn,
        remove_columns=list(processed_train.features),
        batched=True,
        num_proc=4,
        load_from_cache_file=True,
        desc="Tokenizing"
    )
    ...

tokenizer를 클래스의 다른 함수에서도 사용하기 때문에 해당 함수 자체에서 초기화할 수는 없고, 초기화한 tokenizer를 함수 안에서 재선언 해준다.
이렇게 하면 self.tokenizer와 달리 tokenizer는 해당 함수 내에서 최상위에 선언된 객체가 되기 때문에 pickle화 즉, 병렬화가 가능하다.