Untitled
unknown
plain_text
a year ago
27 kB
21
Indexable
# ==============================================================================
# SCRIPT: run_deduplication_rensa_checkpointed.py
# This script uses the high-performance `rensa` library with checkpointing support
# ==============================================================================
import os
import json
import pickle
import pyarrow.parquet as pq
import pandas as pd
import pyarrow as pa
from tqdm import tqdm
import time
from rensa import RMinHash, RMinHashLSH
from pathlib import Path
# ==============================================================================
# 1. CONFIGURATION
# ==============================================================================
BASE_PATH = "D:/pretrain v1"
CHECKPOINT_DIR = f"{BASE_PATH}/dedup_checkpoints"
CONFIG = {
"wikipedia-2k-high-quality": {
"path": f"{BASE_PATH}/deduplicate1/wikipedia-2k-high-quality",
"format": "jsonl",
"text_column": "text",
"priority": 0
},
"arxiv": {
"path": f"{BASE_PATH}/deduplicate1/arxiv",
"format": "jsonl",
"text_column": "content",
"priority": 1
},
"book": {
"path": f"{BASE_PATH}/deduplicate1/book",
"format": "jsonl",
"text_column": "content",
"priority": 2
},
"github": {
"path": f"{BASE_PATH}/deduplicate1/github",
"format": "jsonl",
"text_column": "content",
"priority": 3
},
"stackexchange": {
"path": f"{BASE_PATH}/deduplicate1/stackexchange",
"format": "jsonl", # or "parquet" - check the actual format
"text_column": "content", # or "text" - check the actual column name
"priority": 4 # High priority for technical Q&A content
},
"python-finest-pretrain": {
"path": f"{BASE_PATH}/deduplicate1/python-finest-pretrain/data",
"format": "parquet",
"text_column": "text",
"priority": 5
},
"cc-math-finest": {
"path": f"{BASE_PATH}/deduplicate1/cc-math-finest",
"format": "parquet",
"text_column": "text",
"priority": 6
},
"fineweb-edu-highest-quality-2025": {
"path": f"{BASE_PATH}/deduplicate1/fineweb-edu-highest-quality-2025/data",
"format": "parquet",
"text_column": "text",
"priority": 7
},
"c4": {
"path": f"{BASE_PATH}/deduplicate1/c4",
"format": "jsonl",
"text_column": "content",
"priority": 8
},
"commoncrawl": {
"path": f"{BASE_PATH}/deduplicate1/commoncrawl",
"format": "jsonl",
"text_column": "content",
"priority": 9
}
}
INPUT_DATASET_ROOT = f"{BASE_PATH}/deduplicate1"
FINAL_DATASET_ROOT = f"{BASE_PATH}/final_dataset2"
DEDUP_THRESHOLD = 0.85
NUM_PERM = 128
NGRAM_SIZE = 5
SEED = 42
NUM_BANDS = 16
# ==============================================================================
# 2. CHECKPOINTING UTILITIES
# ==============================================================================
def save_checkpoint(checkpoint_name, data):
"""Save checkpoint data to disk."""
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
checkpoint_path = os.path.join(CHECKPOINT_DIR, f"{checkpoint_name}.pkl")
with open(checkpoint_path, 'wb') as f:
pickle.dump(data, f)
print(f"[CHECKPOINT] Saved: {checkpoint_name}")
def load_checkpoint(checkpoint_name):
"""Load checkpoint data from disk."""
checkpoint_path = os.path.join(CHECKPOINT_DIR, f"{checkpoint_name}.pkl")
if os.path.exists(checkpoint_path):
with open(checkpoint_path, 'rb') as f:
data = pickle.load(f)
print(f"[CHECKPOINT] Loaded: {checkpoint_name}")
return data
return None
def load_checkpoint_safe(checkpoint_name):
"""Load checkpoint data from disk with error handling."""
checkpoint_path = os.path.join(CHECKPOINT_DIR, f"{checkpoint_name}.pkl")
if os.path.exists(checkpoint_path):
try:
with open(checkpoint_path, 'rb') as f:
data = pickle.load(f)
print(f"[CHECKPOINT] Loaded: {checkpoint_name}")
return data
except (EOFError, PermissionError, pickle.UnpicklingError) as e:
print(f"[CHECKPOINT] ERROR loading {checkpoint_name}: {e}")
print(f"[CHECKPOINT] Skipping corrupted checkpoint: {checkpoint_name}")
return None
return None
def get_completed_datasets_safe():
"""Get list of datasets that have been completed, skipping corrupted ones."""
if not os.path.exists(CHECKPOINT_DIR):
return set()
completed = set()
for file in os.listdir(CHECKPOINT_DIR):
if file.startswith("dataset_") and file.endswith(".pkl"):
dataset_name = file.replace("dataset_", "").replace(".pkl", "")
# Try to load it to verify it's not corrupted
if load_checkpoint_safe(f"dataset_{dataset_name}") is not None:
completed.add(dataset_name)
return completed
def process_dataset_with_checkpoint_safe(name, details, lsh, minhashes, doc_metas, idx_to_doc_id, doc_id_to_idx):
"""Process a single dataset with safe checkpointing."""
# Check if this dataset was already processed
checkpoint_data = load_checkpoint_safe(f"dataset_{name}")
if checkpoint_data is not None:
print(f"[RESUME] Dataset '{name}' already processed, loading from checkpoint...")
try:
# Merge checkpoint data back into main structures
minhashes.update(checkpoint_data['minhashes'])
doc_metas.update(checkpoint_data['doc_metas'])
# Update indices
for doc_id in checkpoint_data['doc_ids']:
if doc_id not in doc_id_to_idx:
doc_idx = len(idx_to_doc_id)
idx_to_doc_id.append(doc_id)
doc_id_to_idx[doc_id] = doc_idx
lsh.insert(doc_idx, minhashes[doc_id])
print(f"[RESUME] Loaded {len(checkpoint_data['doc_ids'])} documents from '{name}'")
return len(checkpoint_data['doc_ids'])
except Exception as e:
print(f"[RESUME] Error processing checkpoint for '{name}': {e}")
print(f"[RESUME] Will reprocess '{name}' from scratch")
# Fall through to process fresh
# Process dataset fresh (same as original code)
dataset_path = details["path"]
print(f"\n[PROCESSING] Starting dataset '{name}'...")
if not os.path.exists(dataset_path):
print(f"[DEBUG] ==> ❌ PATH NOT FOUND: {dataset_path}")
return 0
try:
all_files = [f for f in os.listdir(dataset_path) if os.path.isfile(os.path.join(dataset_path, f))]
valid_files = [f for f in all_files if is_valid_data_file(f, details["format"])]
files = sorted([os.path.join(dataset_path, f) for f in valid_files])
except (PermissionError, OSError) as e:
print(f"[DEBUG] ==> ❌ ERROR accessing directory: {e}")
return 0
if not files:
print(f"[DEBUG] ==> ⚠️ No valid {details['format']} files found")
return 0
print(f"[DEBUG] ==> Processing {len(files)} files from '{name}'")
# Store data for this dataset
dataset_minhashes = {}
dataset_doc_metas = {}
dataset_doc_ids = []
docs_processed = 0
for file_path in tqdm(files, desc=f"Reading {name}"):
if details["format"] == "jsonl":
try:
with open(file_path, 'r', encoding='utf-8') as f:
for i, line in enumerate(f):
try:
record = json.loads(line)
text = record.get(details["text_column"], "")
if text:
doc_id = f"{name}::{os.path.basename(file_path)}::{i}"
# Create MinHash
m = RMinHash(num_perm=NUM_PERM, seed=SEED)
shingles = [text[j:j+NGRAM_SIZE] for j in range(len(text) - NGRAM_SIZE + 1)]
if not shingles:
continue
m.update(shingles)
# Store data
doc_idx = len(idx_to_doc_id)
idx_to_doc_id.append(doc_id)
doc_id_to_idx[doc_id] = doc_idx
dataset_minhashes[doc_id] = m
dataset_doc_metas[doc_id] = {"source": name, "priority": details["priority"]}
dataset_doc_ids.append(doc_id)
minhashes[doc_id] = m
doc_metas[doc_id] = dataset_doc_metas[doc_id]
lsh.insert(doc_idx, m)
docs_processed += 1
except (json.JSONDecodeError, AttributeError):
continue
except (PermissionError, OSError) as e:
print(f"[DEBUG] ==> ⚠️ WARNING: Could not read {file_path}: {e}")
continue
elif details["format"] == "parquet":
try:
parquet_file = pq.ParquetFile(file_path)
row_offset = 0
for batch in parquet_file.iter_batches(batch_size=2048):
df = batch.to_pandas()
for idx, row in df.iterrows():
text = row.get(details["text_column"], "")
if text:
doc_id = f"{name}::{os.path.basename(file_path)}::{row_offset + idx}"
# Create MinHash
m = RMinHash(num_perm=NUM_PERM, seed=SEED)
shingles = [text[j:j+NGRAM_SIZE] for j in range(len(text) - NGRAM_SIZE + 1)]
if not shingles:
continue
m.update(shingles)
# Store data
doc_idx = len(idx_to_doc_id)
idx_to_doc_id.append(doc_id)
doc_id_to_idx[doc_id] = doc_idx
dataset_minhashes[doc_id] = m
dataset_doc_metas[doc_id] = {"source": name, "priority": details["priority"]}
dataset_doc_ids.append(doc_id)
minhashes[doc_id] = m
doc_metas[doc_id] = dataset_doc_metas[doc_id]
lsh.insert(doc_idx, m)
docs_processed += 1
row_offset += len(df)
except Exception as e:
print(f"[DEBUG] ==> ⚠️ WARNING: Could not read {file_path}: {e}")
continue
# Try to save checkpoint for this dataset
try:
checkpoint_data = {
'minhashes': dataset_minhashes,
'doc_metas': dataset_doc_metas,
'doc_ids': dataset_doc_ids,
'docs_processed': docs_processed
}
save_checkpoint(f"dataset_{name}", checkpoint_data)
print(f"[CHECKPOINT] Saved progress for '{name}': {docs_processed} documents")
except Exception as e:
print(f"[CHECKPOINT] WARNING: Could not save checkpoint for '{name}': {e}")
print(f"[CHECKPOINT] Continuing without checkpoint for this dataset")
return docs_processed
def get_completed_datasets():
"""Get list of datasets that have been completed."""
if not os.path.exists(CHECKPOINT_DIR):
return set()
completed = set()
for file in os.listdir(CHECKPOINT_DIR):
if file.startswith("dataset_") and file.endswith(".pkl"):
dataset_name = file.replace("dataset_", "").replace(".pkl", "")
completed.add(dataset_name)
return completed
def checkpoint_exists(checkpoint_name):
"""Check if a specific checkpoint exists."""
checkpoint_path = os.path.join(CHECKPOINT_DIR, f"{checkpoint_name}.pkl")
return os.path.exists(checkpoint_path)
# ==============================================================================
# 3. HELPER FUNCTIONS
# ==============================================================================
def is_valid_data_file(file_path, format_type):
"""Check if a file is a valid data file and not a system file."""
filename = os.path.basename(file_path)
if filename.startswith('.') or filename.startswith('_'):
return False
skip_files = {'.cache', '.lock', '.tmp', '.log', 'README', 'LICENSE', '.gitignore',
'.DS_Store', 'Thumbs.db', '.metadata', '.info'}
if filename in skip_files or any(filename.startswith(skip) for skip in skip_files):
return False
if format_type == "jsonl":
return filename.endswith(('.jsonl', '.json', '.ndjson'))
elif format_type == "parquet":
return filename.endswith('.parquet')
return False
# ==============================================================================
# 4. CHECKPOINTED DATA PROCESSING
# ==============================================================================
def process_dataset_with_checkpoint(name, details, lsh, minhashes, doc_metas, idx_to_doc_id, doc_id_to_idx):
"""Process a single dataset with checkpointing."""
# Check if this dataset was already processed
if checkpoint_exists(f"dataset_{name}"):
print(f"[RESUME] Dataset '{name}' already processed, loading from checkpoint...")
checkpoint_data = load_checkpoint_safe(f"dataset_{name}")
# Merge checkpoint data back into main structures
minhashes.update(checkpoint_data['minhashes'])
doc_metas.update(checkpoint_data['doc_metas'])
# Update indices
for doc_id in checkpoint_data['doc_ids']:
if doc_id not in doc_id_to_idx:
doc_idx = len(idx_to_doc_id)
idx_to_doc_id.append(doc_id)
doc_id_to_idx[doc_id] = doc_idx
lsh.insert(doc_idx, minhashes[doc_id])
print(f"[RESUME] Loaded {len(checkpoint_data['doc_ids'])} documents from '{name}'")
return len(checkpoint_data['doc_ids'])
# Process dataset fresh
dataset_path = details["path"]
print(f"\n[PROCESSING] Starting dataset '{name}'...")
if not os.path.exists(dataset_path):
print(f"[DEBUG] ==> ❌ PATH NOT FOUND: {dataset_path}")
return 0
try:
all_files = [f for f in os.listdir(dataset_path) if os.path.isfile(os.path.join(dataset_path, f))]
valid_files = [f for f in all_files if is_valid_data_file(f, details["format"])]
files = sorted([os.path.join(dataset_path, f) for f in valid_files])
except (PermissionError, OSError) as e:
print(f"[DEBUG] ==> ❌ ERROR accessing directory: {e}")
return 0
if not files:
print(f"[DEBUG] ==> ⚠️ No valid {details['format']} files found")
return 0
print(f"[DEBUG] ==> Processing {len(files)} files from '{name}'")
# Store data for this dataset
dataset_minhashes = {}
dataset_doc_metas = {}
dataset_doc_ids = []
docs_processed = 0
for file_path in tqdm(files, desc=f"Reading {name}"):
if details["format"] == "jsonl":
try:
with open(file_path, 'r', encoding='utf-8') as f:
for i, line in enumerate(f):
try:
record = json.loads(line)
text = record.get(details["text_column"], "")
if text:
doc_id = f"{name}::{os.path.basename(file_path)}::{i}"
# Create MinHash
m = RMinHash(num_perm=NUM_PERM, seed=SEED)
shingles = [text[j:j+NGRAM_SIZE] for j in range(len(text) - NGRAM_SIZE + 1)]
if not shingles:
continue
m.update(shingles)
# Store data
doc_idx = len(idx_to_doc_id)
idx_to_doc_id.append(doc_id)
doc_id_to_idx[doc_id] = doc_idx
dataset_minhashes[doc_id] = m
dataset_doc_metas[doc_id] = {"source": name, "priority": details["priority"]}
dataset_doc_ids.append(doc_id)
minhashes[doc_id] = m
doc_metas[doc_id] = dataset_doc_metas[doc_id]
lsh.insert(doc_idx, m)
docs_processed += 1
except (json.JSONDecodeError, AttributeError):
continue
except (PermissionError, OSError) as e:
print(f"[DEBUG] ==> ⚠️ WARNING: Could not read {file_path}: {e}")
continue
elif details["format"] == "parquet":
try:
parquet_file = pq.ParquetFile(file_path)
row_offset = 0
for batch in parquet_file.iter_batches(batch_size=2048):
df = batch.to_pandas()
for idx, row in df.iterrows():
text = row.get(details["text_column"], "")
if text:
doc_id = f"{name}::{os.path.basename(file_path)}::{row_offset + idx}"
# Create MinHash
m = RMinHash(num_perm=NUM_PERM, seed=SEED)
shingles = [text[j:j+NGRAM_SIZE] for j in range(len(text) - NGRAM_SIZE + 1)]
if not shingles:
continue
m.update(shingles)
# Store data
doc_idx = len(idx_to_doc_id)
idx_to_doc_id.append(doc_id)
doc_id_to_idx[doc_id] = doc_idx
dataset_minhashes[doc_id] = m
dataset_doc_metas[doc_id] = {"source": name, "priority": details["priority"]}
dataset_doc_ids.append(doc_id)
minhashes[doc_id] = m
doc_metas[doc_id] = dataset_doc_metas[doc_id]
lsh.insert(doc_idx, m)
docs_processed += 1
row_offset += len(df)
except Exception as e:
print(f"[DEBUG] ==> ⚠️ WARNING: Could not read {file_path}: {e}")
continue
# Save checkpoint for this dataset
checkpoint_data = {
'minhashes': dataset_minhashes,
'doc_metas': dataset_doc_metas,
'doc_ids': dataset_doc_ids,
'docs_processed': docs_processed
}
save_checkpoint(f"dataset_{name}", checkpoint_data)
print(f"[CHECKPOINT] Saved progress for '{name}': {docs_processed} documents")
return docs_processed
# ==============================================================================
# 5. MAIN SCRIPT WITH CHECKPOINTING
# ==============================================================================
def main():
start_time = time.time()
print("=== CHECKPOINTED DEDUPLICATION WITH RESUME CAPABILITY ===")
print(f"Checkpoint directory: {CHECKPOINT_DIR}")
# Check for existing checkpoints
completed_datasets = get_completed_datasets_safe()
if completed_datasets:
print(f"[RESUME] Found checkpoints for: {', '.join(completed_datasets)}")
# --- PHASE 1: Fingerprinting with Checkpointing ---
print("\n--- PHASE 1: Fingerprinting and Indexing Documents ---")
lsh = RMinHashLSH(threshold=DEDUP_THRESHOLD, num_perm=NUM_PERM, num_bands=NUM_BANDS)
minhashes = {}
doc_metas = {}
idx_to_doc_id = []
doc_id_to_idx = {}
total_docs = 0
# Process each dataset with checkpointing
for name, details in CONFIG.items():
docs_processed = process_dataset_with_checkpoint_safe(
name, details, lsh, minhashes, doc_metas, idx_to_doc_id, doc_id_to_idx
)
total_docs += docs_processed
print(f"[PROGRESS] Total documents processed so far: {total_docs}")
if not minhashes:
print("No documents were fingerprinted. Exiting.")
return
# Save main checkpoint after all datasets
main_checkpoint = {
'lsh_state': lsh,
'minhashes': minhashes,
'doc_metas': doc_metas,
'idx_to_doc_id': idx_to_doc_id,
'doc_id_to_idx': doc_id_to_idx,
'total_docs': total_docs
}
save_checkpoint("main_fingerprinting_complete", main_checkpoint)
# --- PHASE 2: Finding and Resolving Duplicate Clusters ---
print("\n--- PHASE 2: Finding and Resolving Duplicate Clusters ---")
# Check if clustering was already done
cluster_checkpoint = load_checkpoint_safe("clustering_complete")
if cluster_checkpoint:
print("[RESUME] Loading clustering results from checkpoint...")
discard_ids = cluster_checkpoint['discard_ids']
clusters = cluster_checkpoint['clusters']
else:
processed_ids = set()
clusters = []
for doc_id in tqdm(minhashes.keys(), desc="Querying for clusters..."):
if doc_id in processed_ids:
continue
result_indices = lsh.query(minhashes[doc_id])
result_ids = [idx_to_doc_id[i] for i in result_indices]
if len(result_ids) > 1:
clusters.append(result_ids)
for member_id in result_ids:
processed_ids.add(member_id)
print(f"INFO: Found {len(clusters)} duplicate clusters.")
discard_ids = set()
for cluster in tqdm(clusters, desc="Resolving clusters"):
best_doc_id = min(cluster, key=lambda doc_id: doc_metas.get(doc_id, {}).get('priority', 99))
for doc_id in cluster:
if doc_id != best_doc_id:
discard_ids.add(doc_id)
# Save clustering checkpoint
cluster_checkpoint_data = {
'discard_ids': discard_ids,
'clusters': clusters
}
save_checkpoint("clustering_complete", cluster_checkpoint_data)
print(f"INFO: Identified {len(discard_ids)} total documents to discard.")
# --- PHASE 3: Writing the Clean Dataset ---
print("\n--- PHASE 3: Writing clean datasets while preserving structure ---")
os.makedirs(FINAL_DATASET_ROOT, exist_ok=True)
total_docs_kept = 0
for name, details in CONFIG.items():
input_dir = details["path"]
if not os.path.exists(input_dir):
continue
# Check if this dataset's output was already processed
if checkpoint_exists(f"output_{name}"):
output_checkpoint = load_checkpoint_safe(f"output_{name}")
total_docs_kept += output_checkpoint['docs_kept']
print(f"[RESUME] Output for '{name}' already exists: {output_checkpoint['docs_kept']} docs kept")
continue
relative_path = os.path.relpath(input_dir, INPUT_DATASET_ROOT)
output_dir = os.path.join(FINAL_DATASET_ROOT, relative_path)
os.makedirs(output_dir, exist_ok=True)
try:
all_files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
valid_files = [f for f in all_files if is_valid_data_file(f, details["format"])]
input_files = sorted([os.path.join(input_dir, f) for f in valid_files])
except (PermissionError, OSError) as e:
print(f"[DEBUG] ==> ❌ ERROR accessing {input_dir}: {e}")
continue
dataset_docs_kept = 0
for input_file_path in tqdm(input_files, desc=f"Filtering {name}"):
output_file_path = os.path.join(output_dir, os.path.basename(input_file_path))
kept_records_count = 0
if details["format"] == "jsonl":
try:
with open(output_file_path, 'w', encoding='utf-8') as f_out:
with open(input_file_path, 'r', encoding='utf-8') as f_in:
for i, line in enumerate(f_in):
doc_id = f"{name}::{os.path.basename(input_file_path)}::{i}"
if doc_id not in discard_ids:
f_out.write(line)
kept_records_count += 1
except Exception as e:
print(f"[DEBUG] ==> ⚠️ WARNING: Could not process {input_file_path}: {e}")
continue
elif details["format"] == "parquet":
try:
df = pd.read_parquet(input_file_path)
ids = [f"{name}::{os.path.basename(input_file_path)}::{i}" for i in range(len(df))]
mask = [doc_id not in discard_ids for doc_id in ids]
filtered_df = df[mask]
if not filtered_df.empty:
pq.write_table(pa.Table.from_pandas(filtered_df, preserve_index=False), output_file_path, compression='snappy')
kept_records_count = len(filtered_df)
except Exception as e:
print(f"[DEBUG] ==> ⚠️ WARNING: Could not process {input_file_path}: {e}")
continue
dataset_docs_kept += kept_records_count
# Save output checkpoint for this dataset
save_checkpoint(f"output_{name}", {'docs_kept': dataset_docs_kept})
total_docs_kept += dataset_docs_kept
end_time = time.time()
print("\n--- Processing Complete! ---")
print(f"Total time taken: {(end_time - start_time) / 60:.2f} minutes")
print(f"Total documents kept: {total_docs_kept}")
print(f"Total documents discarded: {len(discard_ids)}")
print(f"Final structured dataset saved to: '{FINAL_DATASET_ROOT}/'")
print(f"Checkpoints saved to: '{CHECKPOINT_DIR}/'")
if __name__ == "__main__":
main()
Editor is loading...
Leave a Comment