基于Unsloth微调，如何在Hugging Face中无内存加载转换数据集？

我用Unsloth微调Hugging Face模型时，需要将数据集转换为指定格式，为此编写了如下转换函数：

def convert_to_conversation(sample):
    instruction = "OCR the image into markdown format"

    # Get raw image bytes directly
    img_bytes = sample["image"]

    # Convert bytes to base64 string
    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
    img_data_uri = f"data:image/png;base64,{img_b64}"

    return {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": instruction},
                    {"type": "image", "image": img_data_uri}
                ]
            },
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": sample["markdown"]}
                ]
            },
        ]
    }

但使用datasets.map应用该函数时：

mapped_dataset = dataset.map(
    convert_to_conversation,
    remove_columns=dataset.column_names,  # This removes all original columns
    batched=False
)

得到的输出结构异常（出现image: None的无效字段），而用列表推导式处理又会触发内存不足问题。请问如何在不将全量数据集加载到内存的前提下完成转换？

1. 修复结构异常问题 #

输出出现image: None，核心原因是sample["image"]大概率不是原始字节，而是PIL Image对象（Hugging Face Datasets加载图像时默认会转为PIL格式）。修改转换函数，先把PIL图像转成字节再处理：

import io
import base64

def convert_to_conversation(sample):
    instruction = "OCR the image into markdown format"

    # 将PIL Image转为PNG字节
    img = sample["image"]
    img_byte_arr = io.BytesIO()
    img.save(img_byte_arr, format='PNG')
    img_bytes = img_byte_arr.getvalue()

    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
    img_data_uri = f"data:image/png;base64,{img_b64}"

    return {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": instruction},
                    {"type": "image", "image": img_data_uri}
                ]
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": sample["markdown"]}]
            }
        ]
    }

2. 避免内存不足的优化 #

datasets.map本身就是懒加载处理，不会一次性把全量数据加载到内存，只需做好以下几点：

• 保留默认的load_from_cache_file=True，缓存处理后的结果，避免重复计算占用内存
• 开启多进程处理，分散单进程内存压力：

mapped_dataset = dataset.map(
    convert_to_conversation,
    remove_columns=dataset.column_names,
    batched=False,
    num_proc=4,  # 根据自身CPU核心数调整
    load_from_cache_file=True
)

• 若数据集超大，可启用流式加载（仅支持部分数据集格式），处理后保存到磁盘：

from datasets import load_dataset
# 流式加载数据集
dataset = load_dataset("your_dataset_name", streaming=True)

# 流式处理后保存到本地磁盘
mapped_dataset = dataset.map(convert_to_conversation, remove_columns=dataset.column_names)
mapped_dataset.save_to_disk("processed_dataset")

3. 验证处理结果 #

取单个样本检查结构是否正确：

print(mapped_dataset[0]["messages"])

确认user的content同时包含有效文本和图像URI，assistant的文本内容正常。

内容的提问来源于stack exchange，提问作者SoraHeart