File size: 3,270 Bytes
94a2309 a68d8ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import os
import shutil
from datasets import load_from_disk, concatenate_datasets
def combine_datasets(main_dataset_path, datasets_to_add):
"""
Объединяет указанные датасеты с основным датасетом.
Args:
main_dataset_path (str): Путь к основному датасету, в который нужно добавить данные
datasets_to_add (list): Список путей к датасетам, которые нужно добавить
Returns:
Dataset: Объединенный датасет
"""
# Загружаем основной датасет
try:
main_dataset = load_from_disk(main_dataset_path)
print(f"Загружен основной датасет: {main_dataset_path} ({len(main_dataset)} записей)")
except Exception as e:
print(f"Ошибка загрузки основного датасета: {e}")
return None
# Список всех датасетов для объединения
all_datasets = [main_dataset]
# Загружаем и добавляем все дополнительные датасеты
for path in datasets_to_add:
try:
ds = load_from_disk(path)
all_datasets.append(ds)
print(f"Добавлен датасет: {path} ({len(ds)} записей)")
except Exception as e:
print(f"Ошибка загрузки датасета {path}: {e}")
# Объединяем все датасеты
print(f"Объединение {len(all_datasets)} датасетов...")
combined = concatenate_datasets(all_datasets)
# Создаем временную директорию на основе имени основного датасета
temp_dir = f"{main_dataset_path}_temp"
# Удаляем временную директорию, если она уже существует
if os.path.exists(temp_dir):
shutil.rmtree(temp_dir)
try:
# Сохраняем в временную директорию
print(f"Сохранение во временную директорию {temp_dir}...")
combined.save_to_disk(temp_dir)
# Удаляем старую директорию и перемещаем новую на ее место
print(f"Обновление основного датасета...")
#if os.path.exists(main_dataset_path):
# shutil.rmtree(main_dataset_path)
#shutil.copytree(temp_dir, main_dataset_path)
# Удаляем временную директорию после успешного копирования
#shutil.rmtree(temp_dir)
print(f"✅ Объединенный датасет ({len(combined)} записей) успешно сохранен в: {main_dataset_path}")
except Exception as e:
print(f"Ошибка при сохранении датасета: {e}")
print(f"Временные данные сохранены в: {temp_dir}")
return combined
combine_datasets("/workspace/sdxs/datasets/576", ["/workspace/sdxs/datasets/eshu_576"]) |