Issue with datasets.map Freezing during applying a Chat Template

James_Johnson: <p>Hi,</p>
<p>I’m encountering a  issue with the <code>datasets.map</code> function from the Hugging Face <code>datasets</code> library. When I try to apply a custom function to my dataset using <code>map</code>, the process seems to freeze, and I’m unsure why this is happening. This issue does not occur when I run the script  an individually,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/d/7dc617c0bf9ed416e8381b51671228da4616e394.png" data-download-href="https://discuss.pytorch.org/uploads/default/7dc617c0bf9ed416e8381b51671228da4616e394" title="Screenshot from 2024-03-03 06-26-37"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7dc617c0bf9ed416e8381b51671228da4616e394_2_690x350.png" alt="Screenshot from 2024-03-03 06-26-37" data-base62-sha1="hWE5VDLAjNyR26BfjZUwXNEPqja" width="690" height="350" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7dc617c0bf9ed416e8381b51671228da4616e394_2_690x350.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7dc617c0bf9ed416e8381b51671228da4616e394_2_1035x525.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/d/7dc617c0bf9ed416e8381b51671228da4616e394.png 2x" data-dominant-color="505353"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-03-03 06-26-37</span><span class="informations">1109×564 119 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>From the picture, if you notice the percentage remains at 0% even after couple of minutes. I am using ultra_chat dataset. Below are the scripts</p>
<pre><code class="lang-auto"># apply_chat_template.py
import re 
import random 
from multiprocessing import cpu_count 
from load_tokenizer import tokenizer
from load_dataset import raw_datasets
import os
import torch


def apply_chat_template(example,tokenizer):
    print('process example')
    messages = example['messages']
    if messages[0]['role'] != 'system':
        messages.insert(0,{'role':'system','content':''})
    example['text'] = tokenizer.apply_chat_template(messages,tokenize=False)
    return example

print(raw_datasets['train'])

column_names = list(raw_datasets['train'].features)
raw_datasets = raw_datasets.map(apply_chat_template,
                                num_proc=os.cpu_count(),
                                fn_kwargs={'tokenizer':tokenizer},
                                remove_columns=column_names,
                                desc='Applying chat template'
                            )

train_dataset = raw_datasets['train']
test_dataset = raw_datasets['test']
print(train_dataset[0])
</code></pre>
<pre><code class="lang-auto">#load tokenizer.py
from transformers import AutoTokenizer

model_id = 'mistralai/Mistral-7B-v0.1'

tokenizer = AutoTokenizer.from_pretrained(model_id)

# set pad token id to eos token id if pad token id is not set 
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

if tokenizer.model_max_length &gt; 100_000:
    tokenizer.max_model_length = 2048

# set the chat template 
DEFAULT_CHAT_TEMPLATE = "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '&lt;|user|&gt;\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '&lt;|system|&gt;\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '&lt;|assistant|&gt;\n' + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '&lt;|assistant|&gt;'}}\n{% endif %}\n{% endfor %}"
tokenizer.chat_template = DEFAULT_CHAT_TEMPLATE

'''
Below is the structure for chat template 

{% for message in messages %}\n
{% if message['role'] == 'user' %}\n
{{ '&lt;|user|&gt;\n' + message['content'] + eos_token }}\n
{% elif message['role'] == 'system' %}\n
{{ '&lt;|system|&gt;\n' + message['content'] + eos_token }}\n
{% elif message['role'] == 'assistant' %}\n
{{ '&lt;|assistant|&gt;\n' + message['content'] + eos_token }}\n
{% endif %}\n
{% if loop.last and add_generation_prompt %}\n
{{ '&lt;|assistant|&gt;'}}\n
{% endif %}\n
{% endfor %}

'''
'''
#load_dataset.py
from datasets import load_dataset
from datasets import DatasetDict


# load the ultrachat dataset 
ultra_dataset = load_dataset("HuggingFaceH4/ultrachat_200k")

indices = range(0,10)

dataset_dict = {
        'train':ultra_dataset['train_sft'].select(indices=indices),
        'test':ultra_dataset['test_sft'].select(indices=indices)
        }

raw_datasets = DatasetDict(dataset_dict)

# load the first example 
first_example = raw_datasets['train'][0]

# load the messages 
messages = first_example['messages']
# messages is the list of dictionary 

# The length of raw dataset 
print(len(raw_datasets['train']))

</code></pre>