load_data.py

#!/usr/bin/env python3
from datasets import load_dataset
from transformers import AutoTokenizer
import torch
import os
from tqdm import tqdm
import multiprocessing
from functools import partial

NUM_PROC_BASE = max(1, os.cpu_count() // 2 if os.cpu_count() else 1)

TARGET_TOKENS_PER_LANGUAGE = 100_000_000

def tokenize_function(examples, tokenizer):
    output = tokenizer(
        examples["text"],
        add_special_tokens=False,
        truncation=False,
        padding=False,
    )
    return {"input_ids": output.input_ids}

def build_and_save(
    lang="zh",
    date="20231101",
    model_identifier="llama3.1-8b",
    num_proc_map=NUM_PROC_BASE
):
    print(f"Starting data processing for language: {lang}")
    
    tokenizer_name = "meta-llama/Llama-3.1-8B-Instruct" 
    output_dir = "LLama3.1/train-data"
    
    train_filename_base = f"id.{lang}.train.{model_identifier}"
    train_output_path = os.path.join(output_dir, f"{train_filename_base}")

    try:
        ds = load_dataset("wikimedia/wikipedia", f"{date}.{lang}", split="train", trust_remote_code=True)
        if len(ds) == 0:
             print(f"Warning: Dataset for {lang} is empty. Skipping.")
             return
    except Exception as e:
        print(f"Error loading dataset for {lang}: {e}")
        raise

    try:
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_name,
            use_fast=True,
            trust_remote_code=True,
        )
    except Exception as e:
        print(f"Error loading tokenizer '{tokenizer_name}': {e}")
        raise
    
    tokenization_func_with_tokenizer = partial(tokenize_function, tokenizer=tokenizer)
    
    tokenized_ds = ds.map(
        tokenization_func_with_tokenizer,
        batched=True, 
        num_proc=num_proc_map,
        remove_columns=ds.column_names, 
        desc=f"Tokenizing {lang}"
    )
    
    all_document_token_lists = []
    for processed_example in tqdm(tokenized_ds, desc=f"Collecting token lists for {lang}"):
        token_list_for_one_doc = processed_example['input_ids']
        if isinstance(token_list_for_one_doc, list):
            all_document_token_lists.append(token_list_for_one_doc)
            
    if not all_document_token_lists:
        print(f"Warning: No token sequences found for {lang} after tokenization. Skipping.")
        return

    final_token_ids = []
    collected_tokens_count = 0
    for doc_tokens_list in tqdm(all_document_token_lists, desc=f"Aggregating tokens for {lang}"):
        if not doc_tokens_list:
            continue

        current_doc_token_count = len(doc_tokens_list)
        
        if collected_tokens_count + current_doc_token_count <= TARGET_TOKENS_PER_LANGUAGE:
            final_token_ids.extend(doc_tokens_list)
            collected_tokens_count += current_doc_token_count
        else:
            remaining_needed = TARGET_TOKENS_PER_LANGUAGE - collected_tokens_count
            final_token_ids.extend(doc_tokens_list[:remaining_needed])
            collected_tokens_count += remaining_needed
            break 

        if collected_tokens_count >= TARGET_TOKENS_PER_LANGUAGE:
            break
            
    del all_document_token_lists 
    del tokenized_ds
    del ds

    if collected_tokens_count == 0:
        print(f"Warning: Zero tokens collected for {lang}. Skipping save.")
        return

    if collected_tokens_count < TARGET_TOKENS_PER_LANGUAGE:
        print(f"Warning: Language {lang} has only {collected_tokens_count:,} tokens, "
              f"which is less than the target of {TARGET_TOKENS_PER_LANGUAGE:,}.")

    full_tensor = torch.tensor(final_token_ids, dtype=torch.long)
    del final_token_ids

    os.makedirs(output_dir, exist_ok=True)
    torch.save(full_tensor, train_output_path) 
    print(f"Saved {full_tensor.numel():,} tokens for {lang}.")
    del full_tensor


def run_job(args):
    lang, date, model_id, num_proc_map = args
    print(f"Processing language: {lang} (PID: {os.getpid()})")
    try:
        build_and_save(
            lang=lang,
            date=date,
            model_identifier=model_id,
            num_proc_map=num_proc_map
        )
        return lang, True, None 
    except Exception as e:
        import traceback
        traceback.print_exc()
        return lang, False, str(e) 


if __name__ == "__main__":
    languages = ['en', 'zh', 'fr', 'es', 'vi', 'id', 'ja']
    date_snapshot = "20231101" 
    model_id_for_filename = "llama3.1-8b"

    MAX_CONCURRENT_LANGUAGES = 6
    NUM_MAP_PROC_PER_LANG = max(1, NUM_PROC_BASE // MAX_CONCURRENT_LANGUAGES if MAX_CONCURRENT_LANGUAGES > 0 else NUM_PROC_BASE)

    print(f"Starting batch processing for {len(languages)} languages.")

    job_args_list = [
        (lang, date_snapshot, model_id_for_filename, NUM_MAP_PROC_PER_LANG)
        for lang in languages
    ]

    successful_langs = []
    failed_langs_with_errors = {}

    with multiprocessing.Pool(processes=MAX_CONCURRENT_LANGUAGES) as pool:
        results_iterable = pool.imap_unordered(run_job, job_args_list)
        for result in tqdm(results_iterable, total=len(languages), desc="Overall Language Progress"):
            lang_processed, success, error_msg = result
            if success:
                successful_langs.append(lang_processed)
            else:
                failed_langs_with_errors[lang_processed] = error_msg

    print("Batch processing finished.")
    print(f"Successfully processed: {', '.join(sorted(successful_langs))}")
    if failed_langs_with_errors:
        print(f"Failed to process: {', '.join(sorted(failed_langs_with_errors.keys()))}")
        for lang_failed, err in failed_langs_with_errors.items():
            print(f"  - {lang_failed}: {err}")