The datasets.map() implementation modifies the datatype of os.environ object

[leleamol]

In our testing, we noticed that the datasets.map() implementation is modifying the datatype of python os.environ object from '_Environ' to 'dict'.

This causes following function calls to fail as follows:

x = os.environ.get("TEST_ENV_VARIABLE_AFTER_dataset_map", default=None) TypeError: get() takes no keyword arguments
It looks like the following line in datasets.map implementation introduced this functionality.

datasets/src/datasets/arrow_dataset.py

Line 1421 in 0cb1ac0

os.environ = prev_env

Here is the test script to reproduce this error.

from datasets import load_dataset
from transformers import AutoTokenizer
import os


def test_train():
    model_checkpoint = "distilgpt2"
    datasets = load_dataset('wikitext', 'wikitext-2-raw-v1')
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token


    def tokenize_function(examples):
        y = tokenizer(examples['text'], truncation=True, max_length=64)
        return y

    x = os.environ.get("TEST_ENV_VARIABLE_BEFORE_dataset_map", default=None)
    print(f"Testing environment variable: TEST_ENV_VARIABLE_BEFORE_dataset_map {x}")
    print(f"Data type of os.environ before datasets.map = {os.environ.__class__.__name__}")
    datasets.map(tokenize_function, batched=True, num_proc=2, remove_columns=["text"])
    print(f"Data type of os.environ after datasets.map = {os.environ.__class__.__name__}")
    x = os.environ.get("TEST_ENV_VARIABLE_AFTER_dataset_map", default=None)
    print(f"Testing environment variable: TEST_ENV_VARIABLE_AFTER_dataset_map {x}")


if __name__ == "__main__":
    test_train()