script_small.py

import hashlib
import multiprocessing
import os
import random
import re
import sys
import time
from datasets import load_dataset
from huggingface_hub import hf_hub_download
import orjson as json
from itertools import islice
import pickle
import tiktoken
from tqdm import tqdm
from html2term import printc

def load_lille_tokenizer_from_hub(repo_id: str = "Nikity/lille-130m-instruct"):
    """
    Downloads and loads the tiktoken tokenizer for the Lille model from Hugging Face Hub.
    """
    tokenizer_path = hf_hub_download(repo_id=repo_id, filename="tokenizer/Hastings.pkl")
    with open(tokenizer_path, 'rb') as f:
        hastings = pickle.load(f)
    return tiktoken.core.Encoding(hastings.pop('name'), **hastings)

NUM_PROC = os.cpu_count() or 4
BATCH_SIZE = 512

FINAL_DIR = "FINAL"
OUTPUT_FILENAME = os.path.join(FINAL_DIR, "train.txt")
DATA_JSON_FILENAME = os.path.join(FINAL_DIR, "data.json")

MAX_TOKENS = 512
MULTI_TURN_ONLY = False

DATASET_CONFIGS = [
    {"id": "ultrachat_200k", "hf_name": "HuggingFaceH4/ultrachat_200k", "splits": ["train_sft"], "column_map": {"conversation": "messages", "role": "role", "content": "content"}, "roles": {"user": "user", "assistant": "assistant"}},
    {"id": "smoltalk2", "hf_name": ("HuggingFaceTB/smoltalk2", "SFT"), "splits": ["smoltalk_smollm3_smol_magpie_ultra_no_think", "smoltalk_smollm3_everyday_conversations_no_think", "OpenHermes_2.5_no_think", "smoltalk_smollm3_explore_instruct_rewriting_no_think", "smoltalk_smollm3_systemchats_30k_no_think"], "column_map": {"conversation": "messages", "role": "role", "content": "content"}, "roles": {"user": "user", "assistant": "assistant"}},
    {"id": "smol-smoltalk", "hf_name": "HuggingFaceTB/smol-smoltalk", "splits": ["train"], "column_map": {"conversation": "messages", "role": "role", "content": "content"}, "roles": {"user": "user", "assistant": "assistant"}},
    {"id": "WildChat-1M", "hf_name": "allenai/WildChat-1M", "splits": ["train"], "column_map": {"conversation": "conversation", "role": "role", "content": "content"}, "roles": {"user": "user", "assistant": "assistant"}},
    {"id": "ifeval-like-data", "hf_name": ("argilla/ifeval-like-data", "filtered"), "splits": ["train"], "format": "flat", "column_map": {"prompt": "prompt", "response": "response"}},
    {"id": "mmlu", "hf_name": ("cais/mmlu", "all"), "splits": ["auxiliary_train"], "format": "mcq", "column_map": {"question": "question", "choices": "choices", "answer_index": "answer"}},
    {"id": "gsm8k", "hf_name": ("openai/gsm8k", "main"), "splits": ["train"], "format": "flat", "column_map": {"prompt": "question", "response": "answer"}},
    {"id": "tulu-3-sft-personas", "hf_name": "allenai/tulu-3-sft-personas-instruction-following", "splits": ["train"], "column_map": {"conversation": "messages", "role": "role", "content": "content"}, "roles": {"user": "user", "assistant": "assistant"}},
    {"id": "math_qa", "hf_name": "allenai/math_qa", "splits": ["train"], "format": "mcq", "column_map": {"question": "Problem", "choices": "options", "answer_index": "correct"}},
    {"id": "MetaMathQA", "hf_name": "meta-math/MetaMathQA", "splits": ["train"], "format": "flat", "column_map": {"prompt": "query", "response": "response"}},
    {"id": "WizardLM_evol_instruct_V2", "hf_name": "WizardLMTeam/WizardLM_evol_instruct_V2_196k", "splits": ["train"], "column_map": {"conversation": "conversations", "role": "from", "content": "value"}, "roles": {"user": "human", "assistant": "gpt"}},
]

DATASET_TYPE_MAP = {
    "smol-smoltalk": "General Purpose",
    "smoltalk2": "General Purpose",
    "WildChat-1M": "General Purpose",
    "ultrachat_200k": "General Purpose",
    "WizardLM_evol_instruct_V2": "General Purpose",
    "ifeval-like-data": "Instruction",
    "tulu-3-sft-personas": "Instruction",
    "mmlu": "Knowledge",
    "gsm8k": "Math",
    "math_qa": "Math",
    "MetaMathQA": "Math",
}

START_OF_TEXT = "<|startoftext|>"
USER_START = "<|user|>"
ASSISTANT_START = "<|assistant|>"
END_OF_TEXT = "<|endoftext|>"

try:
    import xxhash
    def fast_hash_bytes(b: bytes) -> bytes:
        return xxhash.xxh3_64_digest(b)
except ImportError:
    def fast_hash_bytes(b: bytes) -> bytes:
        return hashlib.blake2b(b, digest_size=16).digest()

TOKENIZER = None

def initialize_tokenizer_for_worker():
    """Initializes the simple_ai tokenizer in each worker process."""
    global TOKENIZER
    if TOKENIZER is None:
        try:
            TOKENIZER = load_lille_tokenizer_from_hub()
        except Exception as e:
            printc(f"🚨 <red>Failed to initialize simple_ai tokenizer in worker {os.getpid()}: {e}</red>", file=sys.stderr)
            sys.exit(1)


def process_batch_worker(args):
    """
    This function runs in a separate process.
    It takes a batch of raw examples, processes them, and returns the results.
    It can handle three formats: 'conversational' (default), 'flat', and 'mcq'.
    """
    global TOKENIZER 

    batch, config = args
    stats = { "processed_raw": 0, "written": 0, "multi_turn": 0, "skipped_format": 0, "skipped_not_multi_turn": 0, "skipped_token_limit": 0, "skipped_encoding_error": 0, "truncated_and_kept": 0 }
    results = []

    if config.get("format") == "flat":
        if MULTI_TURN_ONLY:
            stats["processed_raw"] = len(batch)
            stats["skipped_not_multi_turn"] = len(batch)
            return [], stats
            
        prompt_key = config["column_map"]["prompt"]
        response_key = config["column_map"]["response"]
        
        for example in batch:
            stats["processed_raw"] += 1
            prompt = example.get(prompt_key)
            response = example.get(response_key)

            if not prompt or not isinstance(prompt, str) or not response or not isinstance(response, str):
                stats["skipped_format"] += 1; continue
            
            prompt, response = prompt.strip(), response.strip()
            full_entry = f"{START_OF_TEXT}{USER_START}{prompt}{ASSISTANT_START}{response}{END_OF_TEXT}"
            
            try:
                tokens = TOKENIZER.encode(full_entry, allowed_special="all")
            except Exception:
                stats["skipped_encoding_error"] += 1; continue

            if len(tokens) > MAX_TOKENS:
                stats["skipped_token_limit"] += 1; continue
            
            results.append(full_entry)
            stats["written"] += 1
        return results, stats

    elif config.get("format") == "mcq":
        if MULTI_TURN_ONLY:
            stats["processed_raw"] = len(batch)
            stats["skipped_not_multi_turn"] = len(batch)
            return [], stats
        
        question_key = config["column_map"]["question"]
        choices_key = config["column_map"]["choices"]
        answer_index_key = config["column_map"]["answer_index"]
        
        label_styles = [
            [f"{chr(65+i)}. " for i in range(26)],
            [f"{chr(65+i)}: " for i in range(26)],
            [f"{i+1}. " for i in range(26)],
            [f"({chr(97+i)}) " for i in range(26)],
        ]

        for example in batch:
            stats["processed_raw"] += 1
            question = example.get(question_key)
            choices_raw = example.get(choices_key)
            answer_raw = example.get(answer_index_key)

            choices, answer_index = None, None

            if isinstance(choices_raw, str):
                try:
                    cleaned_choices = [s.strip().strip(',') for s in re.split(r'[a-z]\s*\)\s*', choices_raw) if s.strip()]
                    if cleaned_choices:
                        choices = cleaned_choices
                except Exception:
                    pass
            else:
                choices = choices_raw

            if isinstance(answer_raw, str):
                try:
                    answer_index = ord(answer_raw.strip().lower()) - ord('a')
                except (TypeError, ValueError):
                    pass
            else:
                answer_index = answer_raw
            
            if not all([question, isinstance(question, str), choices, isinstance(choices, list), len(choices) > 1, isinstance(answer_index, int), 0 <= answer_index < len(choices)]):
                stats["skipped_format"] += 1; continue
            
            labels = random.choice(label_styles)
            formatted_choices = "\n".join([f"{labels[i]}{choice}" for i, choice in enumerate(choices)])
            prompt = f"{question.strip()}\n{formatted_choices}"
            response = choices[answer_index].strip()
            full_entry = f"{START_OF_TEXT}{USER_START}{prompt}{ASSISTANT_START}{response}{END_OF_TEXT}"
            
            try:
                tokens = TOKENIZER.encode(full_entry, allowed_special="all")
            except Exception:
                stats["skipped_encoding_error"] += 1; continue

            if len(tokens) > MAX_TOKENS:
                stats["skipped_token_limit"] += 1; continue
            
            results.append(full_entry)
            stats["written"] += 1
        return results, stats

    column_map = config['column_map']
    conv_key, role_key, content_key = column_map["conversation"], column_map["role"], column_map["content"]
    roles_map = config.get("roles", {})
    user_role, assistant_role = roles_map.get("user", "user"), roles_map.get("assistant", "assistant")

    for example in batch:
        stats["processed_raw"] += 1
        conversation_data = example.get(conv_key)
        
        if not isinstance(conversation_data, list) or not conversation_data or \
           conversation_data[0].get(role_key) != user_role or \
           conversation_data[-1].get(role_key) != assistant_role:
            stats["skipped_format"] += 1; continue
            
        string_turns = []
        total_tokens = 2
        
        try:
            for msg in conversation_data:
                role, content = msg.get(role_key), msg.get(content_key, "").strip()
                if not content: continue
                
                prefix = USER_START if role == user_role else (ASSISTANT_START if role == assistant_role else None)
                if not prefix: continue
                
                full_turn_string = f"{prefix}{content}"
                
                turn_tokens = TOKENIZER.encode(full_turn_string, allowed_special="all")
                if total_tokens + len(turn_tokens) > MAX_TOKENS:
                    if string_turns:
                        stats["truncated_and_kept"] += 1
                    break 
                
                string_turns.append(full_turn_string)
                total_tokens += len(turn_tokens)

        except Exception:
            stats["skipped_encoding_error"] += 1
            continue

        if not string_turns:
            stats["skipped_token_limit"] += 1
            continue
        
        if len(string_turns) % 2 != 0:
            string_turns.pop()
            if not string_turns:
                stats["skipped_token_limit"] += 1
                continue

        num_pairs = len(string_turns) // 2
        if MULTI_TURN_ONLY and num_pairs < 2:
            stats["skipped_not_multi_turn"] += 1
            continue
            
        full_entry = f"{START_OF_TEXT}{''.join(string_turns)}{END_OF_TEXT}"
        results.append(full_entry)
        stats["written"] += 1
        if num_pairs >= 2:
            stats["multi_turn"] += 1
            
    return results, stats

def batch_generator(iterable, size):
    """Yields successive n-sized chunks from an iterable."""
    it = iter(iterable)
    while True:
        chunk = tuple(islice(it, size))
        if not chunk:
            return
        yield chunk

def main():
    os.makedirs(FINAL_DIR, exist_ok=True)
    if os.path.exists(OUTPUT_FILENAME):
        printc(f"<yellow>Removing existing output file: <i>{OUTPUT_FILENAME}</i></yellow>")
        os.remove(OUTPUT_FILENAME)

    seen_entries = set()
    global_stats = {}
    
    filter_mode = "multi_turn_only" if MULTI_TURN_ONLY else "all_turns"
    printc("<b>Starting dataset processing pipeline...</b>")
    printc(f"  Output file: <i>{OUTPUT_FILENAME}</i>")
    printc(f"  Parallel processes: <b>{NUM_PROC}</b>")
    printc(f"  Filter mode: <b>{filter_mode}</b>")
    printc(f"  Max tokens:  <b>{MAX_TOKENS}</b><br/>")

    start_time = time.time()
    
    with open(OUTPUT_FILENAME, "a", encoding="utf-8") as out_f, \
         multiprocessing.Pool(processes=NUM_PROC, initializer=initialize_tokenizer_for_worker) as pool:
        is_first_write_for_this_run = True

        for config in DATASET_CONFIGS:
            dataset_id = config["id"]
            printc(f"<cyan>--- Processing Dataset: <b>{dataset_id}</b> ---</cyan>")
            global_stats[dataset_id] = { "written_unique": 0, "skipped_duplicate": 0 }

            for split in config["splits"]:
                try:
                    printc(f"  Streaming and processing split: '<i>{split}</i>'...")
                    load_args = config["hf_name"] if isinstance(config["hf_name"], str) else config["hf_name"][0]
                    load_kwargs = {"name": config["hf_name"][1]} if isinstance(config["hf_name"], tuple) else {}
                    
                    dataset = load_dataset(load_args, **load_kwargs, split=split, streaming=True, trust_remote_code=True)

                    total_examples = None
                    try:
                        total_examples = len(dataset)
                    except TypeError:
                        printc("  <yellow>Note:</yellow> Total examples not available for this streaming split. Progress bar will not show ETA.")

                    batches = batch_generator(dataset, BATCH_SIZE)
                    worker_args = ((batch, config) for batch in batches)
                    
                    results_iterator = pool.imap_unordered(process_batch_worker, worker_args)
                    
                    with tqdm(total=total_examples, desc=f"  -> {split}", unit=" rows", dynamic_ncols=True) as pbar:
                        for processed_texts, batch_stats in results_iterator:
                            for key, value in batch_stats.items():
                                global_stats[dataset_id][key] = global_stats[dataset_id].get(key, 0) + value
                            
                            for entry_text in processed_texts:
                                entry_digest = fast_hash_bytes(entry_text.encode("utf-8"))
                                if entry_digest not in seen_entries:
                                    seen_entries.add(entry_digest)
                                    if not is_first_write_for_this_run:
                                        out_f.write("\n")
                                    out_f.write(entry_text)
                                    is_first_write_for_this_run = False
                                    global_stats[dataset_id]["written_unique"] += 1
                                else:
                                    global_stats[dataset_id]["skipped_duplicate"] += 1
                            
                            if pbar is not None and 'processed_raw' in batch_stats:
                                pbar.update(batch_stats['processed_raw'])

                except Exception as e:
                    printc(f"<br/>🚨 <red><b>An error occurred processing {dataset_id}/{split}:</b> {e}</red>", file=sys.stderr)
            
            ds_stats = global_stats[dataset_id]
            printc(f"  <green>Finished {dataset_id}:</green>")
            for key, value in ds_stats.items():
                if key not in ["written", "multi_turn"]:
                    printc(f"    <#cccccc>{key.replace('_', ' ').capitalize():<25}:</#cccccc> <yellow>{value:>8,}</yellow>")
            printc(f"    <b>{'Entries Written (Unique)':<25}: <green>{ds_stats.get('written_unique', 0):>8,}</green></b><br/>")

    end_time = time.time()
    total_written_count = sum(d.get('written_unique', 0) for d in global_stats.values())
    total_multi_turn_count = sum(d.get('multi_turn', 0) for d in global_stats.values())

    printc("<br/><blue>" + "="*50 + "</blue>")
    printc(f"✅ <b><green>PIPELINE COMPLETE</green> in {end_time - start_time:.2f} seconds</b>")
    printc("<blue>" + "="*50 + "</blue>")
    printc(f"   Filtered output saved to '<i>{OUTPUT_FILENAME}</i>'.")
    
    if MULTI_TURN_ONLY:
        printc(f"   <b>Total unique multi-turn entries written: <yellow>{total_written_count:,}</yellow></b>")
    else:
        printc(f"   <b>Total unique entries written: <yellow>{total_written_count:,}</yellow></b>")
        printc(f"   Hypothetical total if only multi-turn: <yellow>{total_multi_turn_count:,}</yellow>")

    type_counts, dataset_counts = {}, {}
    for ds_id, ds_stats in global_stats.items():
        ds_type = DATASET_TYPE_MAP.get(ds_id, "Unknown")
        written_count = ds_stats.get('written_unique', 0)
        type_counts[ds_type] = type_counts.get(ds_type, 0) + written_count
        dataset_counts[ds_id] = {"type": ds_type, "count": written_count}
    
    type_percentages = {k: (v / total_written_count * 100 if total_written_count > 0 else 0.0) for k, v in type_counts.items()}

    final_json_data = {
        "settings": {"filter_mode": filter_mode, "max_tokens": MAX_TOKENS, "num_proc": NUM_PROC},
        "total_lines_written": total_written_count,
        "multi_turn_line_count": total_multi_turn_count,
        "type_counts": type_counts, "type_percentages": type_percentages,
        "dataset_counts": dataset_counts, "detailed_stats": global_stats
    }

    with open(DATA_JSON_FILENAME, "w", encoding="utf-8") as jf:
        if 'orjson' in sys.modules:
            jf.write(json.dumps(final_json_data, option=json.OPT_INDENT_2).decode('utf-8'))
        else:
            json.dump(final_json_data, jf, indent=2)
    
    printc(f"   Detailed stats saved to '<i>{DATA_JSON_FILENAME}</i>'.<br/>")


if __name__ == "__main__":
    multiprocessing.freeze_support()
    main()