pip install 'litdata[extras]'

# 对于 aws s3
pip install "litdata[extra]" s3fs

# 对于 gcloud storage
pip install "litdata[extra]" gcsfs

from torch.utils.data import DataLoader
from litdata import StreamingRawDataset

dataset = StreamingRawDataset("s3://bucket/files/")

# 配合 PyTorch DataLoader 使用
loader = DataLoader(dataset, batch_size=32)
for batch in loader:
    # 每个元素都是原始字节
    pass

from typing import Union
from torch.utils.data import DataLoader
from litdata import StreamingRawDataset
from litdata.raw.indexer import FileMetadata

class SegmentationRawDataset(StreamingRawDataset):
    def setup(self, files: list[FileMetadata]) -> Union[list[FileMetadata], list[list[FileMetadata]]]:
        # TODO：在此处实现您的自定义分组逻辑。
        # 例如，按前缀、扩展名或任何您需要的规则对文件进行分组。
        # 返回一个组列表，其中每个组都是 FileMetadata 的列表。
        # 示例：
        #   return [[image, mask], ...]
        pass

# 初始化自定义数据集
dataset = SegmentationRawDataset("s3://bucket/files/")
loader = DataLoader(dataset, batch_size=32)
for item in loader:
    # 批次中的每个项目都是一对：[image_bytes, mask_bytes]
    pass

# 当数据集文件发生变化时
dataset = StreamingRawDataset("s3://bucket/files/", recompute_index=True)

from litdata import StreamingDataset, StreamingDataLoader

dataset = StreamingDataset('s3://my-bucket/my-data', shuffle=True)
dataloader = StreamingDataLoader(dataset, batch_size=64)

for batch in dataloader:
    process(batch)  # 替换为您的数据处理逻辑

from litdata import StreamingDataset

# 适用于自定义 S3 兼容端点的 boto3 兼容存储选项
storage_options = {
    "endpoint_url": "your_endpoint_url",
    "aws_access_key_id": "your_access_key_id",
    "aws_secret_access_key": "your_secret_access_key",
}

dataset = StreamingDataset('s3://my-bucket/my-data', storage_options=storage_options)

dataset = StreamingDataset('s3://my-bucket/my-data', storage_options=storage_options)

from litdata import StreamingDataset

# 使用自定义缓存目录初始化 StreamingDataset
dataset = StreamingDataset('s3://my-bucket/my-data', cache_dir="/path/to/cache")

pip install "litdata[extra]" huggingface_hub

# 可选：加速高带宽网络上的下载
pip install hf_transfer
export HF_HUB_ENABLE_HF_TRANSFER=1

import litdata as ld

# 定义 Hugging Face 数据集 URI
hf_dataset_uri = "hf://datasets/leonardPKU/clevr_cogen_a_train/data"

# 创建流式加载数据集
dataset = ld.StreamingDataset(hf_dataset_uri)

# 打印第一个样本
print("Sample", dataset[0])

# 使用 StreamingDataLoader 流式加载数据集
dataloader = ld.StreamingDataLoader(dataset, batch_size=4)
for sample in dataloader:
    pass

import litdata as ld

hf_dataset_uri = "hf://datasets/leonardPKU/clevr_cogen_a_train/data"

ld.index_hf_dataset(hf_dataset_uri)

import litdata as ld

hf_dataset_uri = "hf://datasets/open-thoughts/OpenThoughts-114k/data"

ld.index_parquet_dataset(hf_dataset_uri, "hf-index-dir")

import litdata as ld
from litdata.streaming.item_loader import ParquetLoader

hf_dataset_uri = "hf://datasets/open-thoughts/OpenThoughts-114k/data"

dataset = ld.StreamingDataset(hf_dataset_uri, item_loader=ParquetLoader(), index_path="hf-index-dir")

for batch in ld.StreamingDataLoader(dataset, batch_size=4):
  pass

from litdata import StreamingDataset, StreamingDataLoader

# 对于训练数据集，不要忘记启用 shuffle 和 drop_last !!!
train_dataset = StreamingDataset('s3://my-bucket/my-train-data', shuffle=True, drop_last=True)
train_dataloader = StreamingDataLoader(train_dataset, batch_size=64)

for batch in train_dataloader:
    process(batch)  # 替换为您的数据处理逻辑

val_dataset = StreamingDataset('s3://my-bucket/my-val-data', shuffle=False, drop_last=False)
val_dataloader = StreamingDataLoader(val_dataset, batch_size=64)

for batch in val_dataloader:
    process(batch)  # 替换为您的数据处理逻辑

import os
import litdata as ld

# 使用 boto3 从 AWS S3 读取数据
aws_storage_options={
    "aws_access_key_id": os.environ['AWS_ACCESS_KEY_ID'],
    "aws_secret_access_key": os.environ['AWS_SECRET_ACCESS_KEY'],
}
# 您也可以传递会话选项。（仅适用于 boto3）
aws_session_options = {
  "profile_name": os.environ['AWS_PROFILE_NAME'],  # 仅自定义配置文件需要
  "region_name": os.environ['AWS_REGION_NAME'],    # 仅自定义区域需要
}
dataset = ld.StreamingDataset("s3://my-bucket/my-data", storage_options=aws_storage_options, session_options=aws_session_options)

# 使用未签名请求从 AWS S3 读取数据（使用 boto3）
aws_storage_options={
  "config": botocore.config.Config(
        retries={"max_attempts": 1000, "mode": "adaptive"}, # 配置 S3 操作的重试
        signature_version=botocore.UNSIGNED, # 使用未签名请求
  )
}
dataset = ld.StreamingDataset("s3://my-bucket/my-data", storage_options=aws_storage_options)

aws_storage_options={
    "AWS_ACCESS_KEY_ID": os.environ['AWS_ACCESS_KEY_ID'],
    "AWS_SECRET_ACCESS_KEY": os.environ['AWS_SECRET_ACCESS_KEY'],
    "S3_ENDPOINT_URL": os.environ['AWS_ENDPOINT_URL'],  # 仅自定义端点需要
}
dataset = ld.StreamingDataset("s3://my-bucket/my-data", storage_options=aws_storage_options)

dataset = ld.StreamingDataset("s3://my-bucket/my-data", storage_options=aws_storage_options)

# 从 GCS 读取数据
gcp_storage_options={
    "project": os.environ['PROJECT_ID'],
}
dataset = ld.StreamingDataset("gs://my-bucket/my-data", storage_options=gcp_storage_options)

# 从 Azure 读取数据
azure_storage_options={
    "account_url": f"https://{os.environ['AZURE_ACCOUNT_NAME']}.blob.core.windows.net",
    "credential": os.environ['AZURE_ACCOUNT_ACCESS_KEY']
}
dataset = ld.StreamingDataset("azure://my-bucket/my-data", storage_options=azure_storage_options)

import os
import torch
from litdata import StreamingDataset, StreamingDataLoader

dataset = StreamingDataset("s3://my-bucket/my-data", shuffle=True)
dataloader = StreamingDataLoader(dataset, num_workers=os.cpu_count(), batch_size=64)

# 恢复 dataLoader 状态（如果存在）
if os.path.isfile("dataloader_state.pt"):
    state_dict = torch.load("dataloader_state.pt")
    dataloader.load_state_dict(state_dict)

# 迭代数据
for batch_idx, batch in enumerate(dataloader):

    # 每 1000 个批次存储一次状态
    if batch_idx % 1000 == 0:
        torch.save(dataloader.state_dict(), "dataloader_state.pt")

import numpy as np
from PIL import Image
import litdata as ld

def random_images(index):
    fake_images = Image.fromarray(np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8))
    fake_labels = np.random.randint(10)

    data = {"index": index, "image": fake_images, "class": fake_labels}

    return data

if __name__ == "__main__":
    # optimize 函数将数据写入优化后的格式。
    ld.optimize(
        fn=random_images,                   # 应用于每个输入的函数
        inputs=list(range(1000)),           # 函数的输入（这里是一个数字列表）
        output_dir="fast_data",             # 优化后的数据存储于此
        num_workers=4,                      # 同一台机器上的工作进程数
        chunk_bytes="64MB",                 # 每个数据块的大小
        keep_data_ordered=False,            # 使用共享队列来加速过程
    )

from multiprocessing import Process, Queue
from litdata.processing.data_processor import ALL_DONE
import litdata as ld
import time

def yield_numbers():
    for i in range(1000):
        time.sleep(0.01)
        yield (i, i**2)

def data_producer(q: Queue):
    for item in yield_numbers():
        q.put(item)

    q.put(ALL_DONE)  # 用于指示完成的哨兵值

def fn(index):
    return index  # 用于演示的恒等函数

if __name__ == "__main__":
    q = Queue(maxsize=100)

    producer = Process(target=data_producer, args=(q,))
    producer.start()

    ld.optimize(
        fn=fn,                   # 处理每个项目的函数
        queue=q,                 # 👈 从此队列流式加载数据
        output_dir="fast_data",  # 存储优化数据的位置
        num_workers=2,
        chunk_size=100,
        mode="overwrite",
    )

    producer.join()

import json
from pathlib import Path
import zstandard as zstd
from litdata import optimize, TokensLoader
from tokenizer import Tokenizer
from functools import partial

# 1. 定义一个函数，将 jsonl 文件中的文本转换为 tokens
def tokenize_fn(filepath, tokenizer=None):
    with zstd.open(open(filepath, "rb"), "rt", encoding="utf-8") as f:
        for row in f:
            text = json.loads(row)["text"]
            if json.loads(row)["meta"]["redpajama_set_name"] == "RedPajamaGithub":
                continue  # 排除 GitHub 数据，因为它与 starcoder 重叠
            text_ids = tokenizer.encode(text, bos=False, eos=True)
            yield text_ids

if __name__ == "__main__":
    # 2. 生成输入（我们将优化来自 SlimPajama 数据集的所有压缩 json 文件）
    input_dir = "./slimpajama-raw"
    inputs = [str(file) for file in Path(f"{input_dir}/SlimPajama-627B/train").rglob("*.zst")]

    # 3. 将优化后的数据存储在您想要的任何位置 "/teamspace/datasets" 或 "/teamspace/s3_connections"
    outputs = optimize(
        fn=partial(tokenize_fn, tokenizer=Tokenizer(f"{input_dir}/checkpoints/Llama-2-7b-hf")), # 注意：您可以使用 HF tokenizer 或任何其他
        inputs=inputs,
        output_dir="./slimpajama-optimized",
        chunk_size=(2049 * 8012),
        # 这很重要，用于告知 LitData 我们正在编码连续的 1D 数组（tokens）。
        # LitData 会跳过为每个样本存储元数据，例如所有 tokens 连接形成一个大张量。
        item_loader=TokensLoader(),
    )

import os
from litdata import StreamingDataset, StreamingDataLoader, TokensLoader
from tqdm import tqdm

# 加一，因为我们还需要下一个词
dataset = StreamingDataset(
  input_dir=f"./slimpajama-optimized/train",
  item_loader=TokensLoader(block_size=2048 + 1),
  shuffle=True,
  drop_last=True,
)

train_dataloader = StreamingDataLoader(dataset, batch_size=8, pin_memory=True, num_workers=os.cpu_count())

# 迭代 SlimPajama 数据集
for batch in tqdm(train_dataloader):
    pass

from litdata import optimize, StreamingDataset

def should_keep(index) -> bool:
  # 替换为您自己的逻辑
  return index % 2 == 0

def fn(data):
    if should_keep(data):
        yield data

if __name__ == "__main__":
    optimize(
        fn=fn,
        inputs=list(range(1000)),
        output_dir="only_even_index_optimized",
        chunk_bytes="64MB",
        num_workers=1
    )

    dataset = StreamingDataset("only_even_index_optimized")
    data = list(dataset)
    print(data)
    # [0, 2, 4, 6, 8, 10, ..., 992, 994, 996, 998]

from litdata import optimize, StreamingDataset

def fn(data):
    try:
        yield 1 / data
    except:
        pass

if __name__ == "__main__":
    optimize(
        fn=fn,
        inputs=[0, 0, 0, 1, 2, 4, 0],
        output_dir="only_defined_ratio_optimized",
        chunk_bytes="64MB",
        num_workers=1
    )

    dataset = StreamingDataset("only_defined_ratio_optimized")
    data = list(dataset)
    # 0 被过滤掉，因为它们会引发除零错误
    print(data)
    # [1.0, 0.5, 0.25]