Datasets 数据集库

import datasets
print(datasets.__version__)

from datasets import load_dataset

# 加载Hub上的数据集
dataset = load_dataset("imdb")                       # 返回 DatasetDict
dataset = load_dataset("imdb", split="train")        # 返回 Dataset
dataset = load_dataset("glue", "mrpc", split="train") # 加载子集

# 加载特定版本
dataset = load_dataset("imdb", revision="main")

# 加载特定文件
dataset = load_dataset("username/my_dataset", data_files="train.csv")

# CSV 文件
dataset = load_dataset("csv", data_files="my_data.csv")
dataset = load_dataset("csv", data_files={"train": "train.csv", "test": "test.csv"})

# JSON 文件
dataset = load_dataset("json", data_files="data.json")
dataset = load_dataset("json", data_files="data.jsonl", field="records")  # JSONL格式

# Parquet 文件
dataset = load_dataset("parquet", data_files="data.parquet")

# 文本文件（每行一条）
dataset = load_dataset("text", data_files="sentences.txt")

# 从Python对象直接创建
from datasets import Dataset
dataset = Dataset.from_dict({
    "text": ["Hello", "World"],
    "label": [0, 1]
})
dataset = Dataset.from_pandas(df)      # 从 pandas DataFrame
dataset = Dataset.from_list(records)   # 从字典列表

# 自定义缓存目录
dataset = load_dataset("imdb", cache_dir="/data/cache")

# 设置全局缓存目录
import os
os.environ["HF_DATASETS_CACHE"] = "/data/cache"

from datasets import Dataset

# 基本属性
dataset.column_names     # 列名列表：['text', 'label']
dataset.features         # 特征类型：{'text': Value('string'), 'label': Value('int64')}
dataset.num_rows         # 行数
dataset.num_columns      # 列数
dataset.shape            # (num_rows, num_columns)

# 索引访问
dataset[0]               # 返回第一行的字典：{'text': '...', 'label': 0}
dataset["text"]          # 返回整个text列的列表

# 切片访问
dataset[:5]              # 前五行
dataset.select([0, 2, 4]) # 选择特定行

# 添加/删除列
dataset = dataset.add_column("new_col", [1, 2, 3, ...])
dataset = dataset.remove_columns(["unwanted_col"])
dataset = dataset.rename_column("old_name", "new_name")

from datasets import load_dataset

dataset_dict = load_dataset("imdb")
# DatasetDict({
#     train: Dataset({features: ['text', 'label'], num_rows: 25000})
#     test: Dataset({features: ['text', 'label'], num_rows: 25000})
# })

# 访问分割
train_ds = dataset_dict["train"]
test_ds = dataset_dict["test"]

# 对所有分割统一操作
dataset_dict = dataset_dict.map(tokenize_function, batched=True)
dataset_dict = dataset_dict.filter(lambda x: len(x["text"]) > 10)

# 常用方法
dataset_dict.keys()          # dict_keys(['train', 'test'])
dataset_dict.values()        # 所有Dataset
dataset_dict.items()         # 键值对

def tokenize_function(example):
    return tokenizer(example["text"], truncation=True, padding="max_length", max_length=512)

# 基本用法
dataset = dataset.map(tokenize_function)

# 批量处理（推荐，显著加速）
dataset = dataset.map(tokenize_function, batched=True)

# 多进程并行
dataset = dataset.map(tokenize_function, batched=True, num_proc=8)

# 移除原始列（节省内存）
dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# 只保留指定列
dataset = dataset.map(tokenize_function, batched=True, keep_in_memory=True)

# batched=False 时，函数每次接收一条记录
def process_single(example):
    # example = {"text": "hello", "label": 0}
    example["text_len"] = len(example["text"])
    return example

# batched=True 时，函数每次接收一批记录
def process_batch(examples):
    # examples = {"text": ["hello", "world"], "label": [0, 1]}
    examples["text_len"] = [len(t) for t in examples["text"]]
    return examples

dataset = dataset.map(process_batch, batched=True)

def tokenize_with_max_length(example, max_len=512):
    return tokenizer(example["text"], truncation=True, max_length=max_len)

dataset = dataset.map(
    tokenize_with_max_length,
    fn_kwargs={"max_len": 256},
    batched=True,
)

# 过滤短文本
dataset = dataset.filter(lambda x: len(x["text"]) > 100)

# 过滤特定标签
dataset = dataset.filter(lambda x: x["label"] == 1)

# 使用多进程
dataset = dataset.filter(lambda x: len(x["text"]) > 100, num_proc=4)

# 过滤函数可接受批次
def filter_long_text(examples):
    return [len(t) > 100 for t in examples["text"]]

dataset = dataset.filter(filter_long_text, batched=True)

# 选择指定索引的行
subset = dataset.select([0, 10, 20, 30])

# 选择前1000行
subset = dataset.select(range(1000))

# 随机选择（先shuffle再select）
subset = dataset.shuffle(seed=42).select(range(1000))

# 随机打乱
dataset = dataset.shuffle(seed=42)

# 指定缓冲区大小（大数据集建议增大buffer_size）
dataset = dataset.shuffle(seed=42, buffer_size=10_000)

# 按比例划分
split = dataset.train_test_split(test_size=0.2, seed=42)
# DatasetDict({
#     train: Dataset({...}, num_rows: 800)
#     test: Dataset({...}, num_rows: 200)
# })

# 按绝对数量划分
split = dataset.train_test_split(test_size=1000, seed=42)

# 创建验证集（两次划分）
split1 = dataset.train_test_split(test_size=0.2, seed=42)
split2 = split1["train"].train_test_split(test_size=0.125, seed=42)  # 0.125 * 0.8 = 0.1
# split2["train"] → 72% 训练集
# split2["test"] → 8% 验证集
# split1["test"] → 20% 测试集

# 流式加载（返回 IterableDataset）
dataset = load_dataset("imdb", streaming=True)

# 逐条迭代
for example in dataset["train"]:
    print(example["text"])
    break  # 只看第一条

# 流式数据集的操作
dataset = dataset.shuffle(buffer_size=1000, seed=42)
dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.filter(lambda x: len(x["text"]) > 100)

# 只取前N条
dataset = dataset.take(100)  # 取前100条
dataset = dataset.skip(100)  # 跳过前100条

# 获取少量数据用于调试
small_dataset = list(dataset["train"].take(100))
from datasets import Dataset
debug_dataset = Dataset.from_list(small_dataset)

# 保存为 Arrow 格式
dataset.save_to_disk("my_dataset_path")
dataset_dict.save_to_disk("my_dataset_dict_path")

# 加载
from datasets import load_from_disk
dataset = load_from_disk("my_dataset_path")
dataset_dict = load_from_disk("my_dataset_dict_path")

# 导出为 CSV
dataset.to_csv("output.csv")

# 导出为 Parquet
dataset.to_parquet("output.parquet")

# 导出为 JSON
dataset.to_json("output.jsonl")          # JSONL 格式（每行一条）
dataset.to_json("output.json", orient="records")  # JSON 数组格式

# 推送到Hub
dataset.push_to_hub("username/my_dataset")

# 私有仓库
dataset.push_to_hub("username/my_dataset", private=True)

# 推送 DatasetDict
dataset_dict.push_to_hub("username/my_dataset")

# 设置为 PyTorch 张量格式
dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

# 现在索引访问返回张量
dataset[0]["input_ids"]  # torch.Tensor

# 也可以设置为 numpy
dataset.set_format(type="numpy", columns=["input_ids", "attention_mask", "labels"])

# 也可以设置为 pandas
dataset.set_format(type="pandas")

# 重置格式
dataset.reset_format()

from torch.utils.data import DataLoader

# 方式1：set_format + DataLoader
dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in dataloader:
    input_ids = batch["input_ids"]       # shape: (32, seq_len)
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    # ... 模型前向传播 ...

# 方式2：使用 with_format（不修改原数据集）
torch_dataset = dataset.with_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 1. 加载数据集
dataset = load_dataset("tatsu-lab/alpaca", split="train")

# 2. 定义tokenize函数
def tokenize_function(examples):
    # 将 instruction + input + output 拼接
    prompts = [
        f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}"
        if inp else f"### Instruction:\n{inst}\n\n### Response:\n{out}"
        for inst, inp, out in zip(
            examples["instruction"],
            examples["input"],
            examples["output"]
        )
    ]
    return tokenizer(
        prompts,
        truncation=True,
        max_length=512,
        padding="max_length",
    )

# 3. 批量tokenize
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    num_proc=8,
    remove_columns=dataset.column_names,
)

# 4. 添加labels（对于语言模型，labels = input_ids）
tokenized_dataset = tokenized_dataset.map(
    lambda x: {"labels": x["input_ids"]},
    batched=True,
)

# 5. 划分训练/验证集
split = tokenized_dataset.train_test_split(test_size=0.1, seed=42)

# 6. 设置格式
split.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

# 7. 创建DataLoader
from torch.utils.data import DataLoader
train_loader = DataLoader(split["train"], batch_size=8, shuffle=True)
val_loader = DataLoader(split["test"], batch_size=8)

# 设置下载并发数
dataset = load_dataset("big_dataset", num_proc=8)

dataset = dataset.map(tokenize_function, load_from_cache_file=False)

from datasets import disable_caching, enable_caching

# 全局禁用缓存
disable_caching()

# 重新启用
enable_caching()

from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer.pad_token = tokenizer.eos_token

dataset = load_dataset("tatsu-lab/alpaca", split="train")

def format_and_tokenize(examples):
    # 格式化为对话模板
    texts = []
    for inst, inp, out in zip(examples["instruction"], examples["input"], examples["output"]):
        if inp:
            text = f"Below is an instruction that describes a task, paired with an input.\n\n### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}"
        else:
            text = f"Below is an instruction that describes a task.\n\n### Instruction:\n{inst}\n\n### Response:\n{out}"
        texts.append(text)

    tokenized = tokenizer(texts, truncation=True, max_length=512, padding="max_length")
    tokenized["labels"] = [
        [-100 if t == tokenizer.pad_token_id else t for t in ids]
        for ids in tokenized["input_ids"]
    ]
    return tokenized

processed = dataset.map(
    format_and_tokenize,
    batched=True,
    num_proc=8,
    remove_columns=dataset.column_names,
)
processed.save_to_disk("./alpaca_tokenized")

from datasets import load_dataset

# 流式加载 C4 数据集（约300GB）
c4 = load_dataset("allenai/c4", "en", streaming=True, split="train")

# 定义预处理函数
def preprocess(examples):
    # 只保留长度适中的文本
    return {"text": [t for t in examples["text"] if 100 < len(t) < 10000]}

c4 = c4.filter(lambda x: 100 < len(x["text"]) < 10000)
c4 = c4.shuffle(buffer_size=10_000, seed=42)

# 逐批迭代，不需要将整个数据集加载到内存
for i, example in enumerate(c4):
    if i >= 100000:  # 只处理前10万条
        break
    text = example["text"]
    # ... 处理文本 ...

from datasets import Dataset, DatasetDict
import json

# 从本地文件加载
with open("train.jsonl") as f:
    train_data = [json.loads(line) for line in f]
with open("test.jsonl") as f:
    test_data = [json.loads(line) for line in f]

train_ds = Dataset.from_list(train_data)
test_ds = Dataset.from_list(test_data)

# 创建 DatasetDict
ds_dict = DatasetDict({"train": train_ds, "test": test_ds})

# 推送到 Hub
ds_dict.push_to_hub("my-username/my-llm-dataset", private=True)

from datasets import load_dataset
import re

dataset = load_dataset("openwebtext", split="train", streaming=True)

def is_quality_text(text):
    # 长度过滤
    if len(text) < 200 or len(text) > 100000:
        return False
    # 去重检查（简单hash）
    # 重复行检测
    lines = text.split("\n")
    unique_ratio = len(set(lines)) / max(len(lines), 1)
    if unique_ratio < 0.5:
        return False
    # 语言检测（简单启发式）
    if len(re.findall(r'[\u4e00-\u9fff]', text)) / max(len(text), 1) > 0.3:
        return True  # 中文内容
    return True

dataset = dataset.filter(lambda x: is_quality_text(x["text"]))

# Dataset 的核心属性
dataset._data         # pyarrow.Table，底层Arrow表格
dataset._indices      # 可选的索引映射（用于select/sort等操作的惰性实现）
dataset._format_type  # 输出格式：None/"torch"/"numpy"/"pandas"
dataset._format_columns  # 格式化的列