script_parquet.py

import hashlib
import multiprocessing
import os
import random
import re
import sys
import time
from datasets import load_dataset, 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_PARQUET"
OUTPUT_FILENAME = os.path.join(FINAL_DIR, "train.parquet")
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 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 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 structured results
    suitable for Parquet output.
    It can handle three formats: 'conversational' (default), 'flat', and 'mcq'.
    """
    global TOKENIZER 

    batch, config, hf_dataset = 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_formatted_str = f"{START_OF_TEXT}{USER_START}{prompt}{ASSISTANT_START}{response}{END_OF_TEXT}"
            
            try:
                tokens = TOKENIZER.encode(full_formatted_str, allowed_special="all")
            except Exception:
                stats["skipped_encoding_error"] += 1; continue

            if len(tokens) > MAX_TOKENS:
                stats["skipped_token_limit"] += 1; continue
            
            messages = [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": response}
            ]
            results.append({"messages": messages, "hf_dataset": hf_dataset})
            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_formatted_str = f"{START_OF_TEXT}{USER_START}{prompt}{ASSISTANT_START}{response}{END_OF_TEXT}"
            
            try:
                tokens = TOKENIZER.encode(full_formatted_str, allowed_special="all")
            except Exception:
                stats["skipped_encoding_error"] += 1; continue

            if len(tokens) > MAX_TOKENS:
                stats["skipped_token_limit"] += 1; continue
            
            messages = [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": response}
            ]
            results.append({"messages": messages, "hf_dataset": hf_dataset})
            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_for_token_count = []
        structured_messages = []
        total_tokens_approx = 0

        try:
            for msg_idx, msg in enumerate(conversation_data):
                role, content = msg.get(role_key), msg.get(content_key, "").strip()
                if not content: continue
                
                prefix = None
                hf_role = None
                if role == user_role:
                    prefix = USER_START
                    hf_role = "user"
                elif role == assistant_role:
                    prefix = ASSISTANT_START
                    hf_role = "assistant"
                else:
                    continue
                
                full_turn_string_for_token_count = f"{prefix}{content}"
                
                current_turn_tokens = TOKENIZER.encode(full_turn_string_for_token_count, allowed_special="all")
                
                estimated_total_tokens = total_tokens_approx + len(current_turn_tokens) + (2 if msg_idx == 0 else 0)

                if estimated_total_tokens > MAX_TOKENS:
                    if structured_messages:
                        stats["truncated_and_kept"] += 1
                    break
                
                string_turns_for_token_count.append(full_turn_string_for_token_count)
                structured_messages.append({"role": hf_role, "content": content})
                total_tokens_approx = estimated_total_tokens

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

        if not structured_messages:
            stats["skipped_token_limit"] += 1
            continue
        
        if structured_messages[-1]["role"] == "user":
            structured_messages.pop()
            if not structured_messages:
                stats["skipped_token_limit"] += 1
                continue

        num_pairs = len(structured_messages) // 2
        if MULTI_TURN_ONLY and num_pairs < 2:
            stats["skipped_not_multi_turn"] += 1
            continue

        results.append({"messages": structured_messages, "hf_dataset": hf_dataset})
        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 = {}
    final_dataset_records = []

    filter_mode = "multi_turn_only" if MULTI_TURN_ONLY else "all_turns"
    printc("<b>Starting dataset processing pipeline (Parquet output)...</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 multiprocessing.Pool(processes=NUM_PROC, initializer=initialize_tokenizer_for_worker) as pool:

        for config in DATASET_CONFIGS:
            dataset_id = config["id"]
            hf_name_for_column = config["hf_name"] if isinstance(config["hf_name"], str) else config["hf_name"][0]
            
            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, hf_name_for_column) 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_batch_results, 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 record in processed_batch_results:
                                record_digest = fast_hash_bytes(json.dumps(record["messages"], option=json.OPT_SORT_KEYS))
                                if record_digest not in seen_entries:
                                    seen_entries.add(record_digest)
                                    final_dataset_records.append(record)
                                    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 = len(final_dataset_records)
    if total_written_count > 0:
        printc(f"<blue>Creating Hugging Face Dataset from {total_written_count:,} unique records...</blue>")
        final_hf_dataset = Dataset.from_list(final_dataset_records)
        printc(f"<blue>Saving dataset to Parquet: <i>{OUTPUT_FILENAME}</i></blue>")
        final_hf_dataset.to_parquet(OUTPUT_FILENAME)
    else:
        printc("<yellow>No unique records were processed and collected. Parquet file will not be created.</yellow>")

    total_multi_turn_count_pre_dedupe = 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 to Parquet: <yellow>{total_written_count:,}</yellow></b>")
    else:
        printc(f"   <b>Total unique entries written to Parquet: <yellow>{total_written_count:,}</yellow></b>")
        printc(f"   Hypothetical total *initially processed* if only multi-turn (pre-deduplication): <yellow>{total_multi_turn_count_pre_dedupe:,}</yellow>")

    final_type_counts = {}
    final_dataset_specific_counts = {}
    
    for record in final_dataset_records:
        hf_dataset = record["hf_dataset"]
        
        original_ds_id = next((cfg["id"] for cfg in DATASET_CONFIGS 
                               if (isinstance(cfg["hf_name"], tuple) and cfg["hf_name"][0] == hf_dataset) 
                               or (isinstance(cfg["hf_name"], str) and cfg["hf_name"] == hf_dataset)), "Unknown Dataset")
        
        ds_type = DATASET_TYPE_MAP.get(original_ds_id, "Unknown")
        final_type_counts[ds_type] = final_type_counts.get(ds_type, 0) + 1
        final_dataset_specific_counts[hf_dataset] = final_dataset_specific_counts.get(hf_dataset, 0) + 1

    final_type_percentages = {k: (v / total_written_count * 100 if total_written_count > 0 else 0.0) for k, v in final_type_counts.items()}
    
    final_dataset_counts_by_id = {}
    for cfg in DATASET_CONFIGS:
        cfg_hf_name = cfg["hf_name"] if isinstance(cfg["hf_name"], str) else cfg["hf_name"][0]
        final_dataset_counts_by_id[cfg["id"]] = {
            "type": DATASET_TYPE_MAP.get(cfg["id"], "Unknown"), 
            "count": final_dataset_specific_counts.get(cfg_hf_name, 0)
        }

    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_pre_dedupe": total_multi_turn_count_pre_dedupe, 
        "type_counts": final_type_counts, 
        "type_percentages": final_type_percentages,
        "dataset_counts": final_dataset_counts_by_id, 
        "detailed_stats": global_stats 
    }

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


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