from datasets import load_dataset
dataset = load_dataset("wikitext", "wikitext-103-v1")

train_file = '/mnt/data/vinhbk/wikitext103/train.txt'
valid_file = '/mnt/data/vinhbk/wikitext103/valid.txt'
test_file = '/mnt/data/vinhbk/wikitext103/test.txt'

# Save the train split
with open(train_file, 'w', encoding='utf-8') as f:
    for item in dataset['train']:
        f.write(item['text'] + '\n')

# Save the validation split
with open(valid_file, 'w', encoding='utf-8') as f:
    for item in dataset['validation']:
        f.write(item['text'] + '\n')

# Save the test split
with open(test_file, 'w', encoding='utf-8') as f:
    for item in dataset['test']:
        f.write(item['text'] + '\n')

print("Files saved successfully!")