from datasets import Dataset, DatasetDict

data = {
    "text": [
        "이 논문은 transformer architecture의 scalability를 잘 설명하고 있다",
        "모델의 convergence 속도가 느리고 training instability 문제가 보인다",
        "실험 결과에서 overfitting 현상이 명확하게 관찰된다",
        "proposed method는 baseline 대비 높은 accuracy를 달성했다",
        "loss function 설계가 충분히 justified되지 않았다",
        "data preprocessing 단계가 systematic하게 잘 구성되어 있다",
    ],
    "label": [1, 0, 0, 1, 0, 1],
}

dataset = Dataset.from_dict(data)

dataset_dict = DatasetDict({
    "train": dataset.select([0, 1, 2, 3]),
    "validation": dataset.select([4, 5]),
})

def add_length(example):
    return {"length": len(example["text"])}

dataset_mapped = dataset.map(add_length)
dataset_mapped[0]

def add_length_batch(batch):
    return {"length": [len(t) for t in batch["text"]]}

dataset_mapped = dataset.map(
    add_length_batch,
    batched=True,
    batch_size=2,
)

dataset_dict_mapped = dataset_dict.map(
    add_length_batch,
    batched=True,
)

def keep_long(example):
    return len(example["text"].strip()) >= 10

dataset_filtered = dataset.filter(keep_long)

dataset_dict_filtered = dataset_dict.filter(keep_long)

dataset_shuffled = dataset.shuffle(seed=42)

dataset_dict_shuffled = dataset_dict.shuffle(seed=42)

dataset_sorted = dataset.sort("label")

dataset_selected = dataset.select([0, 3, 5])

dataset_shard0 = dataset.shard(num_shards=2, index=0)
dataset_shard1 = dataset.shard(num_shards=2, index=1)

split = dataset.train_test_split(
    test_size=0.3,
    seed=42,
)

dataset_no_text = dataset.remove_columns(["text"])

dataset_dict_no_text = dataset_dict.remove_columns(["text"])

dataset_renamed = dataset.rename_column("label", "y")

dataset_renamed = dataset.rename_columns({
    "text": "sentence",
    "label": "y",
})

dataset_text_only = dataset.select_columns(["text"])

from datasets import ClassLabel, Features, Value

features = Features({
    "text": Value("string"),
    "label": ClassLabel(names=["neg", "pos"]),
})

dataset_cast = dataset.cast(features)

dataset_torch = dataset.with_format(
    "torch",
    columns=["label"],
)

dataset.set_format("numpy", columns=["label"])
dataset.reset_format()

def to_upper(example):
    example["text"] = example["text"].upper()
    return example

dataset_transformed = dataset.with_transform(to_upper)

from datasets import load_from_disk

dataset.save_to_disk("tmp_dataset")
dataset_loaded = load_from_disk("tmp_dataset")

dataset.to_csv("dataset.csv")
dataset.to_json("dataset.json")
dataset.to_parquet("dataset.parquet")

# dataset.push_to_hub("username/my-dataset")
# dataset_dict.push_to_hub("username/my-dataset-dict")

dataset_dict.keys()
dataset_dict["train"][0]

for split, ds in dataset_dict.items():
    print(split, len(ds))