data/build.py

import os
import json
import re
import argparse
from datetime import datetime
import numpy as np

from data.prepare import clean_and_normalize, get_next_version
from data.tokenize import BPETokenizer

def check_quality(filename, text):
    warnings = []
    # 1. Very short documents
    if len(text.strip()) < 50:
        warnings.append("Extremely short document (less than 50 characters).")
        
    # 2. Mostly whitespace
    whitespace_ratio = len(re.findall(r'\s', text)) / max(1, len(text))
    if whitespace_ratio > 0.5:
        warnings.append(f"Mostly whitespace (whitespace ratio: {whitespace_ratio:.1%}).")
        
    # 3. Mostly symbols
    symbols = re.findall(r'[^a-zA-Z0-9\s]', text)
    symbol_ratio = len(symbols) / max(1, len(text))
    if symbol_ratio > 0.4:
        warnings.append(f"Mostly symbols/punctuation (symbol ratio: {symbol_ratio:.1%}).")
        
    # 4. Extremely repetitive text (check 3-gram frequencies)
    words = re.findall(r'\w+', text.lower())
    if len(words) >= 10:
        trigrams = list(zip(words[:-2], words[1:-1], words[2:]))
        if trigrams:
            freq = {}
            for tg in trigrams:
                freq[tg] = freq.get(tg, 0) + 1
            max_freq = max(freq.values())
            ratio = max_freq / len(trigrams)
            if ratio > 0.2:
                warnings.append(f"Highly repetitive phrasing detected (top trigram frequency ratio: {ratio:.1%}).")

    return warnings

def main():
    parser = argparse.ArgumentParser(description="COLLISION Dataset Builder")
    parser.add_argument("--raw-dir", type=str, default="data/raw", help="Path to raw data directory")
    parser.add_argument("--datasets-dir", type=str, default="datasets", help="Path to datasets directory")
    parser.add_argument("--save-dir", type=str, default="artifacts/tokenizer", help="Tokenizer save directory")
    parser.add_argument("--vocab-size", type=int, default=8000, help="Target vocabulary size")
    parser.add_argument("--seed", type=int, default=42, help="Seed for deterministic data splitting")
    args = parser.parse_args()

    # Discover files
    if not os.path.exists(args.raw_dir):
        print(f"Error: Raw directory not found at {args.raw_dir}")
        return

    txt_files = [f for f in os.listdir(args.raw_dir) if f.endswith(".txt")]
    if not txt_files:
        print(f"No .txt documents discovered in {args.raw_dir}.")
        return

    print(f"--- Discovering raw documents ---")
    documents = []
    seen_hashes = set()
    duplicate_count = 0
    empty_file_count = 0
    total_characters = 0
    quality_warnings = {}

    for f_name in txt_files:
        f_path = os.path.join(args.raw_dir, f_name)
        with open(f_path, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()
            
        if not content.strip():
            empty_file_count += 1
            quality_warnings[f_name] = ["Empty document file."]
            continue

        total_characters += len(content)
        cleaned = clean_and_normalize(content)
        
        # Check suspicious patterns
        warnings = check_quality(f_name, cleaned)
        if warnings:
            quality_warnings[f_name] = warnings

        # Deduplication
        doc_hash = hash(cleaned)
        if doc_hash in seen_hashes:
            duplicate_count += 1
            continue
            
        seen_hashes.add(doc_hash)
        documents.append({
            "name": f_name,
            "text": cleaned
        })

    print(f"Found {len(txt_files)} file(s).")
    print(f"Empty files: {empty_file_count} | Duplicates removed: {duplicate_count}")
    print(f"Unique documents remaining: {len(documents)}")

    if not documents:
        print("Error: No unique content remaining to tokenize.")
        return

    # Train/Reload Tokenizer
    print("\n--- Tokenizer Training & Setup ---")
    combined_clean_text = "\n".join([doc["text"] for doc in documents])
    tokenizer = BPETokenizer()
    training_text_subset = combined_clean_text[:20000]
    stats = tokenizer.train(training_text_subset, args.vocab_size)
    tokenizer.save(args.save_dir)

    
    with open(os.path.join(args.save_dir, "stats.json"), "w", encoding="utf-8") as f:
        json.dump(stats, f, indent=2)
        
    print(f"Tokenizer vocabulary size: {len(tokenizer.inverse_vocab)}")

    # Verificationencode/decode
    test_str = "The future of artificial intelligence."
    enc_test = tokenizer.encode(test_str, bos=True, eos=True)
    dec_test = tokenizer.decode(enc_test)
    is_valid = dec_test.strip() == test_str.strip() or test_str.replace(" ", "") in dec_test.replace(" ", "")
    print(f"Tokenizer Verification | Input: '{test_str}' | Output: '{dec_test.strip()}' | Reconstructed matches: {is_valid}")

    # Deterministic Split
    print("\n--- Deterministic Dataset Splitting (90/10) ---")
    # Set seed locally for split
    rng = np.random.default_rng(args.seed)
    indices = np.arange(len(documents))
    rng.shuffle(indices)
    
    split_idx = int(0.9 * len(documents))
    train_indices = indices[:split_idx]
    val_indices = indices[split_idx:]
    
    # If single document, allocate at least one to val for safety
    if len(documents) == 1:
        # Fallback to token level split
        print("Single document corpus detected. Performing split at token level instead of document level.")
        token_ids = tokenizer.encode(combined_clean_text, bos=True, eos=True)
        split_tok_idx = int(0.9 * len(token_ids))
        train_tokens = token_ids[:split_tok_idx]
        val_tokens = token_ids[split_tok_idx:]
    else:
        train_tokens = []
        for idx in train_indices:
            train_tokens.extend(tokenizer.encode(documents[idx]["text"], bos=True, eos=True))
            
        val_tokens = []
        for idx in val_indices:
            val_tokens.extend(tokenizer.encode(documents[idx]["text"], bos=True, eos=True))

    total_tokens = len(train_tokens) + len(val_tokens)

    # Determine Version Directory
    next_ver = get_next_version(args.datasets_dir)
    version_name = f"collision_dataset_v{next_ver}"
    version_dir = os.path.join(args.datasets_dir, version_name)
    os.makedirs(version_dir, exist_ok=True)

    # Save to Version Dir
    v_train_path = os.path.join(version_dir, "train.bin")
    v_val_path = os.path.join(version_dir, "val.bin")
    np.array(train_tokens, dtype=np.uint16).tofile(v_train_path)
    np.array(val_tokens, dtype=np.uint16).tofile(v_val_path)

    # Copy to global processed dir
    g_train_path = os.path.join("data", "processed", "train.bin")
    g_val_path = os.path.join("data", "processed", "val.bin")
    os.makedirs(os.path.dirname(g_train_path), exist_ok=True)
    np.array(train_tokens, dtype=np.uint16).tofile(g_train_path)
    np.array(val_tokens, dtype=np.uint16).tofile(g_val_path)

    # Print Quality warning summary
    if quality_warnings:
        print("\n--- DATASET QUALITY WARNINGS (Needs Review) ---")
        for doc_name, warns in quality_warnings.items():
            print(f"  [{doc_name}]:")
            for w in warns:
                print(f"    - {w}")

    # Generate metadata.json
    doc_lens = [len(doc["text"]) for doc in documents]
    meta = {
        "dataset_version": version_name,
        "creation_date": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        "source_files": txt_files,
        "number_of_files": len(txt_files),
        "number_of_documents": len(documents),
        "character_count": total_characters,
        "token_count": total_tokens,
        "train_tokens": len(train_tokens),
        "validation_tokens": len(val_tokens),
        "tokenizer_version": "1.0-BPETokenizer",
        "vocabulary_size": len(tokenizer.inverse_vocab),
        "avg_document_length_chars": float(np.mean(doc_lens)) if doc_lens else 0.0,
        "min_document_length_chars": int(np.min(doc_lens)) if doc_lens else 0,
        "max_document_length_chars": int(np.max(doc_lens)) if doc_lens else 0,
        "duplicate_count": duplicate_count,
        "empty_file_count": empty_file_count,
        "preprocessing_configuration": {
            "whitespace_normalized": True,
            "duplicate_filtering": True,
            "split_ratio": "90/10",
            "seed": args.seed
        },
        "quality_warnings": quality_warnings
    }
    
    meta_path = os.path.join(version_dir, "metadata.json")
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)

    print(f"\nSaved dataset version metadata to {meta_path}")
    print(f"Dataset Build completed successfully. Token count: {total_tokens:,}")

if __name__ == "__main__":
    main()