Create prepare.py

ff23d5b verified 1 day ago

9.65 kB

import os
import numpy as np
from tokenizers import Tokenizer
from datasets import load_dataset
from huggingface_hub import HfApi, list_repo_files
import time

# ==========================================
# НАСТРОЙКИ
# ==========================================
HF_TOKEN = "hf_Abc123"
REPO_ID = "ExtraX1/TinyCyclic-0.4B-Base"
api = HfApi()

# (файл ~100 МБ при 50M токенов uint16)
CHUNK_SIZE_TOKENS = 50_000_000

print("📥 Загружаем токенизатор...")
tokenizer = Tokenizer.from_pretrained("ExtraX1/TinyCyclic-Tokenizer")
eot_id = tokenizer.token_to_id("</s>")

def process_domain(domain_name, sources, max_tokens):
    print(f"\n⬇️ Подготовка [{domain_name.upper()}] (цель: {max_tokens:,} токенов)...")

    # Проверяем, что уже есть на HF
    try:
        existing_files = list_repo_files(repo_id=REPO_ID, token=HF_TOKEN)
    except Exception as e:
        print(f"❌ Ошибка доступа к репозиторию: {e}")
        return

    prefix = f"prepared_tokens/{domain_name}/"
    uploaded_chunks = [f for f in existing_files if f.startswith(prefix) and f.endswith(".bin")]

    start_chunk_idx = len(uploaded_chunks)
    already_downloaded_tokens = start_chunk_idx * CHUNK_SIZE_TOKENS

    if already_downloaded_tokens >= max_tokens:
        print(f"✅ Для [{domain_name.upper()}] уже скачано достаточно ({already_downloaded_tokens:,} токенов). Пропускаем.")
        return

    print(f"⏩ Уже загружено чанков: {start_chunk_idx}. Продолжаем...")

    current_tokens = []
    chunk_idx = start_chunk_idx
    total_tokens_this_run = 0

    for source in sources:
        print(f"  📚 Источник: {source['name']}...")
        try:
            load_kwargs = {
                "split": source.get("split", "train"),
                "streaming": True,
            }
            
            # Если это прямой путь к Parquet-файлам (для ZH)
            if source.get("data_files"):
                load_kwargs["path"] = "parquet"
                load_kwargs["data_files"] = source["data_files"]
            else:
                # Для всех остальных (EN, RU, CODE, LOGIC)
                load_kwargs["path"] = source["path"]
                if source.get("config"):
                    load_kwargs["name"] = source["config"]
                    
            ds = load_dataset(**load_kwargs)
        except Exception as e:
            print(f"  ⚠️ Не удалось загрузить {source['name']}: {e}. Пропускаем.")
            continue

        for item in ds:
            if already_downloaded_tokens + total_tokens_this_run >= max_tokens:
                break

            # --- Логика извлечения текста ---
            if source.get("nested"):
                # Stack v3: item — это репозиторий, внутри item["files"] список файлов
                files = item.get("files", [])
                target_langs = source.get("target_languages", [])
                for file in files:
                    lang = file.get("language", "")
                    if target_langs and lang not in target_langs:
                        continue
                    text = file.get("content", "")
                    if not text or not text.strip():
                        continue

                    ids = tokenizer.encode(text).ids
                    current_tokens.extend(ids)
                    current_tokens.append(eot_id)
                    total_tokens_this_run += len(ids)

                    if len(current_tokens) >= CHUNK_SIZE_TOKENS:
                        filename = f"temp_{domain_name}_{chunk_idx:04d}.bin"
                        arr = np.array(current_tokens, dtype=np.uint16)
                        arr.tofile(filename)
                        print(f"  ⬆️ Заливаем {domain_name}/chunk_{chunk_idx:04d}.bin ({len(arr):,} токенов)...")
                        try:
                            api.upload_file(
                                path_or_fileobj=filename,
                                path_in_repo=f"prepared_tokens/{domain_name}/chunk_{chunk_idx:04d}.bin",
                                repo_id=REPO_ID,
                                token=HF_TOKEN
                            )
                            print("    ✅ Успешно.")
                        except Exception as e:
                            print(f"    ❌ Ошибка загрузки: {e}. Ждем 30 сек...")
                            time.sleep(30)
                        os.remove(filename)
                        current_tokens = []
                        chunk_idx += 1

                    if already_downloaded_tokens + total_tokens_this_run >= max_tokens:
                        break
            else:
                # Обычный плоский датасет (EN, RU, ZH, LOGIC)
                text = item.get(source['field'], "")
                if not text or not text.strip():
                    continue

                ids = tokenizer.encode(text).ids
                current_tokens.extend(ids)
                current_tokens.append(eot_id)
                total_tokens_this_run += len(ids)

                if len(current_tokens) >= CHUNK_SIZE_TOKENS:
                    filename = f"temp_{domain_name}_{chunk_idx:04d}.bin"
                    arr = np.array(current_tokens, dtype=np.uint16)
                    arr.tofile(filename)
                    print(f"  ⬆️ Заливаем {domain_name}/chunk_{chunk_idx:04d}.bin ({len(arr):,} токенов)...")
                    try:
                        api.upload_file(
                            path_or_fileobj=filename,
                            path_in_repo=f"prepared_tokens/{domain_name}/chunk_{chunk_idx:04d}.bin",
                            repo_id=REPO_ID,
                            token=HF_TOKEN
                        )
                        print("    ✅ Успешно.")
                    except Exception as e:
                        print(f"    ❌ Ошибка загрузки: {e}. Ждем 30 сек...")
                        time.sleep(30)
                    os.remove(filename)
                    current_tokens = []
                    chunk_idx += 1

        if already_downloaded_tokens + total_tokens_this_run >= max_tokens:
            break

    if current_tokens:
        filename = f"temp_{domain_name}_{chunk_idx:04d}.bin"
        arr = np.array(current_tokens, dtype=np.uint16)
        arr.tofile(filename)
        print(f"  ⬆️ Заливаем последний {domain_name}/chunk_{chunk_idx:04d}.bin ({len(arr):,} токенов)...")
        try:
            api.upload_file(
                path_or_fileobj=filename,
                path_in_repo=f"prepared_tokens/{domain_name}/chunk_{chunk_idx:04d}.bin",
                repo_id=REPO_ID,
                token=HF_TOKEN
            )
            print("    ✅ Успешно.")
        except Exception as e:
            print(f"    ❌ Ошибка загрузки: {e}")
        os.remove(filename)

# ==========================================
# ЗАПУСК (~5B)
# ==========================================

# 1. АНГЛИЙСКИЙ (2.4B токенов)
en_sources = [
    {"name": "FineWeb-Edu", "path": "HuggingFaceFW/fineweb-edu", "config": "sample-10BT", "field": "text"},
    {"name": "Philosophy (EN)", "path": "LisaMegaWatts/philosophy-corpus", "config": None, "field": "text"}
]
process_domain("en", en_sources, max_tokens=2_400_000_000)

# 2. РУССКИЙ (1.4B токенов)
ru_sources = [
    {"name": "Wikipedia RU", "path": "wikimedia/wikipedia", "config": "20231101.ru", "field": "text"},
    {"name": "DaruLM Habr", "path": "dichspace/darulm", "config": "habr", "field": "text"},
    {"name": "DaruLM Textbooks", "path": "dichspace/darulm", "config": "textbook", "field": "text"},
    {"name": "DaruLM Philosophy", "path": "dichspace/darulm", "config": "philosophy", "field": "text"}
]
process_domain("ru", ru_sources, max_tokens=1_400_000_000)

# 3. КОД (0.75B токенов)
# Stack v3: каждый item — репозиторий, код внутри item["files"], фильтруем по language
code_sources = [
    {
        "name": "Stack v3 (Multilingual)",
        "path": "HuggingFaceCode/stack-v3-train",
        "nested": True,
        "target_languages": [
            "Python", "C", "C++", "C#", "Java", "JavaScript",
            "Go", "Rust", "Shell", "Markdown"
        ]
    }
]
process_domain("code", code_sources, max_tokens=750_000_000)

# 4. МАТЕМАТИКА и ЛОГИКА (0.5B токенов)
# MetaMathQA + OpenMathInstruct-2 (чтобы добрать объем)
logic_sources = [
    {"name": "MetaMathQA", "path": "meta-math/MetaMathQA", "config": None, "field": "query"},
    {"name": "OpenMathInstruct-2", "path": "nvidia/OpenMathInstruct-2", "config": None, "field": "problem"}
]
process_domain("logic", logic_sources, max_tokens=500_000_000)

# 5. КИТАЙСКИЙ (0.25B токенов)
# Используем конфиг pretrain (сплит score_4_5 — самый качественный), а не SFT
zh_sources = [
    {
        "name": "Fineweb-Edu-Chinese-V2.2 (Pretrain)",
        "data_files": "hf://datasets/opencsg/Fineweb-Edu-Chinese-V2.2@ab9e282ac7f146b52ea19df00c0225146b784129/4_5/*.parquet",
        "field": "text"
    }
]
process_domain("zh", zh_sources, max_tokens=250_000_000)

print("\n🎉 ВСЁ ГОТОВО!")

	import os
	import numpy as np
	from tokenizers import Tokenizer
	from datasets import load_dataset
	from huggingface_hub import HfApi, list_repo_files
	import time

	# ==========================================
	# НАСТРОЙКИ
	# ==========================================
	HF_TOKEN = "hf_Abc123"
	REPO_ID = "ExtraX1/TinyCyclic-0.4B-Base"
	api = HfApi()

	# (файл ~100 МБ при 50M токенов uint16)
	CHUNK_SIZE_TOKENS = 50_000_000

	print("📥 Загружаем токенизатор...")
	tokenizer = Tokenizer.from_pretrained("ExtraX1/TinyCyclic-Tokenizer")
	eot_id = tokenizer.token_to_id("</s>")

	def process_domain(domain_name, sources, max_tokens):
	print(f"\n⬇️ Подготовка [{domain_name.upper()}] (цель: {max_tokens:,} токенов)...")

	# Проверяем, что уже есть на HF
	try:
	existing_files = list_repo_files(repo_id=REPO_ID, token=HF_TOKEN)
	except Exception as e:
	print(f"❌ Ошибка доступа к репозиторию: {e}")
	return

	prefix = f"prepared_tokens/{domain_name}/"
	uploaded_chunks = [f for f in existing_files if f.startswith(prefix) and f.endswith(".bin")]

	start_chunk_idx = len(uploaded_chunks)
	already_downloaded_tokens = start_chunk_idx * CHUNK_SIZE_TOKENS

	if already_downloaded_tokens >= max_tokens:
	print(f"✅ Для [{domain_name.upper()}] уже скачано достаточно ({already_downloaded_tokens:,} токенов). Пропускаем.")
	return

	print(f"⏩ Уже загружено чанков: {start_chunk_idx}. Продолжаем...")

	current_tokens = []
	chunk_idx = start_chunk_idx
	total_tokens_this_run = 0

	for source in sources:
	print(f" 📚 Источник: {source['name']}...")
	try:
	load_kwargs = {
	"split": source.get("split", "train"),
	"streaming": True,
	}

	# Если это прямой путь к Parquet-файлам (для ZH)
	if source.get("data_files"):
	load_kwargs["path"] = "parquet"
	load_kwargs["data_files"] = source["data_files"]
	else:
	# Для всех остальных (EN, RU, CODE, LOGIC)
	load_kwargs["path"] = source["path"]
	if source.get("config"):
	load_kwargs["name"] = source["config"]

	ds = load_dataset(**load_kwargs)
	except Exception as e:
	print(f" ⚠️ Не удалось загрузить {source['name']}: {e}. Пропускаем.")
	continue

	for item in ds:
	if already_downloaded_tokens + total_tokens_this_run >= max_tokens:
	break

	# --- Логика извлечения текста ---
	if source.get("nested"):
	# Stack v3: item — это репозиторий, внутри item["files"] список файлов
	files = item.get("files", [])
	target_langs = source.get("target_languages", [])
	for file in files:
	lang = file.get("language", "")
	if target_langs and lang not in target_langs:
	continue
	text = file.get("content", "")
	if not text or not text.strip():
	continue

	ids = tokenizer.encode(text).ids
	current_tokens.extend(ids)
	current_tokens.append(eot_id)
	total_tokens_this_run += len(ids)

	if len(current_tokens) >= CHUNK_SIZE_TOKENS:
	filename = f"temp_{domain_name}_{chunk_idx:04d}.bin"
	arr = np.array(current_tokens, dtype=np.uint16)
	arr.tofile(filename)
	print(f" ⬆️ Заливаем {domain_name}/chunk_{chunk_idx:04d}.bin ({len(arr):,} токенов)...")
	try:
	api.upload_file(
	path_or_fileobj=filename,
	path_in_repo=f"prepared_tokens/{domain_name}/chunk_{chunk_idx:04d}.bin",
	repo_id=REPO_ID,
	token=HF_TOKEN
	)
	print(" ✅ Успешно.")
	except Exception as e:
	print(f" ❌ Ошибка загрузки: {e}. Ждем 30 сек...")
	time.sleep(30)
	os.remove(filename)
	current_tokens = []
	chunk_idx += 1

	if already_downloaded_tokens + total_tokens_this_run >= max_tokens:
	break
	else:
	# Обычный плоский датасет (EN, RU, ZH, LOGIC)
	text = item.get(source['field'], "")
	if not text or not text.strip():
	continue

	ids = tokenizer.encode(text).ids
	current_tokens.extend(ids)
	current_tokens.append(eot_id)
	total_tokens_this_run += len(ids)

	if len(current_tokens) >= CHUNK_SIZE_TOKENS:
	filename = f"temp_{domain_name}_{chunk_idx:04d}.bin"
	arr = np.array(current_tokens, dtype=np.uint16)
	arr.tofile(filename)
	print(f" ⬆️ Заливаем {domain_name}/chunk_{chunk_idx:04d}.bin ({len(arr):,} токенов)...")
	try:
	api.upload_file(
	path_or_fileobj=filename,
	path_in_repo=f"prepared_tokens/{domain_name}/chunk_{chunk_idx:04d}.bin",
	repo_id=REPO_ID,
	token=HF_TOKEN
	)
	print(" ✅ Успешно.")
	except Exception as e:
	print(f" ❌ Ошибка загрузки: {e}. Ждем 30 сек...")
	time.sleep(30)
	os.remove(filename)
	current_tokens = []
	chunk_idx += 1

	if already_downloaded_tokens + total_tokens_this_run >= max_tokens:
	break

	if current_tokens:
	filename = f"temp_{domain_name}_{chunk_idx:04d}.bin"
	arr = np.array(current_tokens, dtype=np.uint16)
	arr.tofile(filename)
	print(f" ⬆️ Заливаем последний {domain_name}/chunk_{chunk_idx:04d}.bin ({len(arr):,} токенов)...")
	try:
	api.upload_file(
	path_or_fileobj=filename,
	path_in_repo=f"prepared_tokens/{domain_name}/chunk_{chunk_idx:04d}.bin",
	repo_id=REPO_ID,
	token=HF_TOKEN
	)
	print(" ✅ Успешно.")
	except Exception as e:
	print(f" ❌ Ошибка загрузки: {e}")
	os.remove(filename)

	# ==========================================
	# ЗАПУСК (~5B)
	# ==========================================

	# 1. АНГЛИЙСКИЙ (2.4B токенов)
	en_sources = [
	{"name": "FineWeb-Edu", "path": "HuggingFaceFW/fineweb-edu", "config": "sample-10BT", "field": "text"},
	{"name": "Philosophy (EN)", "path": "LisaMegaWatts/philosophy-corpus", "config": None, "field": "text"}
	]
	process_domain("en", en_sources, max_tokens=2_400_000_000)

	# 2. РУССКИЙ (1.4B токенов)
	ru_sources = [
	{"name": "Wikipedia RU", "path": "wikimedia/wikipedia", "config": "20231101.ru", "field": "text"},
	{"name": "DaruLM Habr", "path": "dichspace/darulm", "config": "habr", "field": "text"},
	{"name": "DaruLM Textbooks", "path": "dichspace/darulm", "config": "textbook", "field": "text"},
	{"name": "DaruLM Philosophy", "path": "dichspace/darulm", "config": "philosophy", "field": "text"}
	]
	process_domain("ru", ru_sources, max_tokens=1_400_000_000)

	# 3. КОД (0.75B токенов)
	# Stack v3: каждый item — репозиторий, код внутри item["files"], фильтруем по language
	code_sources = [
	{
	"name": "Stack v3 (Multilingual)",
	"path": "HuggingFaceCode/stack-v3-train",
	"nested": True,
	"target_languages": [
	"Python", "C", "C++", "C#", "Java", "JavaScript",
	"Go", "Rust", "Shell", "Markdown"
	]
	}
	]
	process_domain("code", code_sources, max_tokens=750_000_000)

	# 4. МАТЕМАТИКА и ЛОГИКА (0.5B токенов)
	# MetaMathQA + OpenMathInstruct-2 (чтобы добрать объем)
	logic_sources = [
	{"name": "MetaMathQA", "path": "meta-math/MetaMathQA", "config": None, "field": "query"},
	{"name": "OpenMathInstruct-2", "path": "nvidia/OpenMathInstruct-2", "config": None, "field": "problem"}
	]
	process_domain("logic", logic_sources, max_tokens=500_000_000)

	# 5. КИТАЙСКИЙ (0.25B токенов)
	# Используем конфиг pretrain (сплит score_4_5 — самый качественный), а не SFT
	zh_sources = [
	{
	"name": "Fineweb-Edu-Chinese-V2.2 (Pretrain)",
	"data_files": "hf://datasets/opencsg/Fineweb-Edu-Chinese-V2.2@ab9e282ac7f146b52ea19df00c0225146b784129/4_5/*.parquet",
	"field": "text"
	}
	]
	process_domain("zh", zh_sources, max_tokens=250_000_000)

	print("\n🎉 ВСЁ ГОТОВО!")