# -*- coding: utf-8 -*- """ DEDUPLICATE litdata_pretrain_final — remove near-duplicate documents. Strategy: 1. Scan all 308 chunks, hash each document (first 200 tokens) 2. Keep first occurrence, mark subsequent duplicates for removal 3. Rebuild chunks with duplicates removed (same format, new files) 4. Update index.json Memory-efficient: processes 10 chunks at a time, uses hash set. """ import json import os import time import hashlib from pathlib import Path import numpy as np ROOT = Path(__file__).resolve().parent.parent.parent FINAL_DIR = ROOT / "Base" / "data" / "litdata_pretrain_final" BLOCK_SIZE = 1025 DTYPE = np.int32 EOS_TOKEN_ID = 0 CHUNK_BYTES_TARGET = 64 * 1024 * 1024 HASH_WINDOW = 200 # first N tokens for hash def read_chunk(filepath): with open(filepath, "rb") as f: raw = f.read() num_blocks = np.frombuffer(raw[:4], dtype=np.uint32)[0] header_size = 4 + (num_blocks + 1) * 4 data_bytes = raw[header_size:] expected_tokens = num_blocks * BLOCK_SIZE expected_bytes = expected_tokens * DTYPE().itemsize tokens = np.frombuffer(data_bytes[:expected_bytes], dtype=DTYPE) return tokens, int(num_blocks) def extract_documents(tokens): """Extract individual documents separated by EOS.""" eos_positions = np.where(tokens == EOS_TOKEN_ID)[0] docs = [] start = 0 for eos_pos in eos_positions: if eos_pos > start: docs.append(tokens[start:eos_pos]) start = eos_pos + 1 # Trailing partial (no EOS at end — crosses chunk boundary) if start < len(tokens): remaining = tokens[start:] if len(remaining) > 0: docs.append(remaining) return docs def doc_hash(token_array): """Hash first HASH_WINDOW tokens of a document.""" key = token_array[:HASH_WINDOW].tobytes() return hashlib.md5(key).hexdigest() def write_chunk(filepath, tokens_array, block_size=BLOCK_SIZE): """Write a litdata chunk from a flat token array.""" num_blocks = len(tokens_array) // block_size if num_blocks == 0: return None actual = num_blocks * block_size data = np.array(tokens_array[:actual], dtype=DTYPE) header_num = np.array([num_blocks], dtype=np.uint32) offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (block_size * DTYPE().itemsize) header = np.concatenate([header_num, offsets]) with open(filepath, "wb") as f: header.tofile(f) data.tofile(f) return { "chunk_bytes": int(header.nbytes + data.nbytes), "chunk_size": num_blocks, "dim": int(actual), "filename": os.path.basename(filepath), } class StreamingDeduplicator: """Accumulates deduplicated tokens and writes chunks.""" def __init__(self, output_dir, backup_suffix="_dedup"): self.output_dir = Path(output_dir) self.dtype_size = DTYPE().itemsize self.tokens_per_chunk = (CHUNK_BYTES_TARGET // self.dtype_size // BLOCK_SIZE) * BLOCK_SIZE self.buffer = [] self.chunk_idx = 0 self.chunks_meta = [] self.total_tokens = 0 def add_doc(self, doc_tokens): self.buffer.extend(doc_tokens.tolist()) self.buffer.append(EOS_TOKEN_ID) while len(self.buffer) >= self.tokens_per_chunk: self._flush() def _flush(self): if len(self.buffer) < BLOCK_SIZE: return take = min(len(self.buffer), self.tokens_per_chunk) num_blocks = take // BLOCK_SIZE if num_blocks == 0: return actual = num_blocks * BLOCK_SIZE data = np.array(self.buffer[:actual], dtype=DTYPE) self.buffer = self.buffer[actual:] filename = f"chunk-0-{self.chunk_idx}.bin" filepath = self.output_dir / filename header_num = np.array([num_blocks], dtype=np.uint32) offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * self.dtype_size) header = np.concatenate([header_num, offsets]) with open(filepath, "wb") as f: header.tofile(f) data.tofile(f) meta = { "chunk_bytes": int(header.nbytes + data.nbytes), "chunk_size": num_blocks, "dim": int(actual), "filename": filename, } self.chunks_meta.append(meta) self.total_tokens += actual self.chunk_idx += 1 if self.chunk_idx % 25 == 0: print(f" Written {self.chunk_idx} deduped chunks ({self.total_tokens:,} tokens)") def finalize(self): while len(self.buffer) >= BLOCK_SIZE: self._flush() discarded = len(self.buffer) self.buffer = [] return self.total_tokens, discarded def main(): t_start = time.time() with open(FINAL_DIR / "index.json") as f: index = json.load(f) chunks_meta = index["chunks"] config = index.get("config", {}) num_chunks = len(chunks_meta) original_tokens = sum(c["dim"] for c in chunks_meta) print(f"{'='*75}") print(f" DEDUPLICATING litdata_pretrain_final") print(f"{'='*75}") print(f" Chunks: {num_chunks}") print(f" Original tokens: {original_tokens:,}") print() # ── PASS 1: Scan all chunks, collect hashes, identify duplicates ────── print(f" PASS 1: Scanning all chunks for duplicates...") seen_hashes = set() dup_count = 0 keep_count = 0 total_docs = 0 # We need to track which is first occurrence # But since we process sequentially, just check "in seen_hashes" # Use a temp dir to write deduplicated data, then swap TEMP_DIR = FINAL_DIR.parent / "litdata_pretrain_dedup_temp" if TEMP_DIR.exists(): import shutil shutil.rmtree(str(TEMP_DIR)) os.makedirs(str(TEMP_DIR)) writer = StreamingDeduplicator(TEMP_DIR) for ci, meta in enumerate(chunks_meta): filepath = FINAL_DIR / meta["filename"] tokens, num_blocks = read_chunk(filepath) docs = extract_documents(tokens) chunk_dups = 0 chunk_kept = 0 for doc in docs: total_docs += 1 if len(doc) < 10: # Very short fragments — keep (usually chunk-boundary partials) writer.add_doc(doc) keep_count += 1 chunk_kept += 1 continue h = doc_hash(doc) if h in seen_hashes: dup_count += 1 chunk_dups += 1 else: seen_hashes.add(h) writer.add_doc(doc) keep_count += 1 chunk_kept += 1 if (ci + 1) % 25 == 0 or ci == num_chunks - 1: print(f" Chunk {ci+1}/{num_chunks}: total docs={total_docs:,}, kept={keep_count:,}, dupes removed={dup_count:,}") # Finalize final_tokens, discarded = writer.finalize() new_chunks = writer.chunk_idx print(f"\n PASS 1 COMPLETE:") print(f" Total documents scanned: {total_docs:,}") print(f" Documents kept: {keep_count:,}") print(f" Duplicates removed: {dup_count:,} ({100*dup_count/max(total_docs,1):.2f}%)") print(f" Unique hashes: {len(seen_hashes):,}") print(f" Tokens after dedup: {final_tokens:,}") print(f" Token reduction: {original_tokens - final_tokens:,} ({100*(original_tokens-final_tokens)/original_tokens:.2f}%)") print(f" Chunks after dedup: {new_chunks}") print(f" Discarded partial: {discarded} tokens") # ── PASS 2: Swap temp into final ────────────────────────────────────── print(f"\n PASS 2: Replacing original with deduplicated data...") # Remove old chunk files for meta in chunks_meta: old_file = FINAL_DIR / meta["filename"] if old_file.exists(): os.remove(str(old_file)) # Move new chunk files from temp to final import shutil for meta in writer.chunks_meta: src = TEMP_DIR / meta["filename"] dst = FINAL_DIR / meta["filename"] shutil.move(str(src), str(dst)) # Remove temp dir shutil.rmtree(str(TEMP_DIR)) # Update index.json new_index = { "chunks": writer.chunks_meta, "config": config, "updated_at": str(time.time()), } with open(FINAL_DIR / "index.json", "w") as f: json.dump(new_index, f, indent=2) elapsed = time.time() - t_start # ── Report ──────────────────────────────────────────────────────────── report = [] report.append(f"{'='*75}") report.append(f" DEDUPLICATION REPORT — litdata_pretrain_final") report.append(f"{'='*75}") report.append(f"") report.append(f" Time: {elapsed:.0f}s ({elapsed/60:.1f} min)") report.append(f"") report.append(f" BEFORE:") report.append(f" Chunks: {num_chunks}") report.append(f" Tokens: {original_tokens:,}") report.append(f" Documents: {total_docs:,}") report.append(f"") report.append(f" AFTER:") report.append(f" Chunks: {new_chunks}") report.append(f" Tokens: {final_tokens:,}") report.append(f" Documents: {keep_count:,}") report.append(f"") report.append(f" REMOVED:") report.append(f" Duplicate docs: {dup_count:,} ({100*dup_count/max(total_docs,1):.2f}%)") report.append(f" Tokens removed: {original_tokens - final_tokens:,} ({100*(original_tokens-final_tokens)/original_tokens:.2f}%)") report.append(f"") report.append(f" Format: litdata binary (int32, BLOCK_SIZE={BLOCK_SIZE}, EOS={EOS_TOKEN_ID})") report.append(f" Location: {FINAL_DIR}") report.append(f"{'='*75}") full_report = '\n'.join(report) print(f"\n{full_report}") with open(FINAL_DIR / "DEDUP_REPORT.txt", "w", encoding="utf-8") as f: f.write(full_report) print(f"\n Saved to: {FINAL_DIR / 'DEDUP_REPORT.txt'}") print(f" Done! Dataset is now clean and deduplicated.") if __name__ == "__main__": main()