added coderabbit changes

Author: DhruvK278

openverifiablellm/config.py

@@ -29,21 +29,7 @@ class BenchmarkConfig:
 
 
 def _load_from_yaml(yaml_path: Path) -> Optional[List[str]]:
-    """
-    Attempt to read benchmark identifiers from a YAML file.
-
-    Expected format::
-
-        benchmarks:
-          - gsm8k
-          - cais/mmlu
-          - /path/to/local/dataset.jsonl
-
-    Returns:
-        Optional[List[str]]: A list of benchmark identifiers (e.g., "gsm8k", "cais/mmlu", or local file paths)
-        when the YAML is present and valid, or None when the file is missing, unreadable, or does not contain
-        a "benchmarks" list.
-    """
+    """ Attempt to read benchmark identifiers from a YAML file. """
     if not yaml_path.is_file():
         return None
 
@@ -68,7 +54,6 @@ def load_config(cli_benchmarks: Optional[str] = None) -> BenchmarkConfig:
     """Build a :class:`BenchmarkConfig` by resolving benchmark sources."""
     config = BenchmarkConfig()
 
-    # Priority 1: CLI flag
     if cli_benchmarks:
         benchmarks = [b.strip() for b in cli_benchmarks.split(",") if b.strip()]
         if benchmarks:
@@ -78,13 +63,11 @@ def load_config(cli_benchmarks: Optional[str] = None) -> BenchmarkConfig:
         else:
             logger.warning("CLI provided empty benchmarks; falling back to next priority.")
 
-    # Priority 2: YAML config file
     yaml_benchmarks = _load_from_yaml(Path(BENCHMARKS_YAML))
     if yaml_benchmarks:
         config.benchmarks = yaml_benchmarks
         logger.info("Benchmarks from %s: %s", BENCHMARKS_YAML, config.benchmarks)
         return config
 
-    # Priority 3: defaults
     logger.info("Using default benchmarks: %s", config.benchmarks)
     return config

openverifiablellm/contamination.py

@@ -14,7 +14,7 @@
 import logging
 import re
 from pathlib import Path
-from typing import Iterable, List, Optional, Set
+from typing import Dict, List, Optional, Set
 
 from rbloom import Bloom
 
@@ -58,7 +58,11 @@ class DatasetLoadError(Exception):
     pass
 
 
-def _load_hf(dataset_id: str, trust_remote_code: bool = False) -> List[str]:
+def _load_hf(
+    dataset_id: str,
+    trust_remote_code: bool = False,
+    split_map: Optional[Dict[str, List[str]]] = None
+) -> List[str]:
     """Load benchmark texts from a Hugging Face dataset."""
     try:
         from datasets import load_dataset
@@ -68,14 +72,25 @@ def _load_hf(dataset_id: str, trust_remote_code: bool = False) -> List[str]:
             "Install it with: pip install datasets"
         ) from exc
 
+    if split_map is None:
+        split_map = {"gsm8k": ["test"],"cais/mmlu": ["test", "validation"]}
+
     texts: List[str] = []
+    splits_to_load = split_map.get(dataset_id)
+    
     try:
-        ds = load_dataset(dataset_id, trust_remote_code=trust_remote_code)
+        if splits_to_load:
+            splits = [
+                load_dataset(dataset_id, split=s, trust_remote_code=trust_remote_code)
+                for s in splits_to_load
+            ]
+        else:
+            ds = load_dataset(dataset_id, trust_remote_code=trust_remote_code)
+            splits = ds.values() if hasattr(ds, "values") else [ds]
     except Exception as exc:
         logger.error("Could not load HF dataset '%s': %s", dataset_id, exc)
         raise DatasetLoadError(f"Failed to load HF dataset '{dataset_id}'") from exc
 
-    splits = ds.values() if hasattr(ds, "values") else [ds]
     for split in splits:
         for row in split:
             for field in _TEXT_FIELDS:
@@ -163,7 +178,9 @@ def build_bloom_filter(
 
     hasher = hashlib.sha256()
     for text in benchmark_texts:
-        hasher.update(text.encode("utf-8"))
+        encoded = text.encode("utf-8")
+        hasher.update(len(encoded).to_bytes(8, "little"))
+        hasher.update(encoded)
     inputs_hash = hasher.hexdigest()
 
     current_meta = {
@@ -229,8 +246,9 @@ def build_bloom_filter(
 def check_contamination(
     text: str,
     bloom_filter: Bloom,
-    benchmark_texts: List[str],
+    benchmark_texts: Optional[List[str]] = None,
     n: int = 13,
+    precomputed_normalised_benchmarks: Optional[Set[str]] = None,
 ) -> bool:
     """
     Determine whether text is contaminated by benchmark data.
@@ -244,18 +262,19 @@ def check_contamination(
        performed to eliminate false positives.
 
     """
+    if benchmark_texts is None and precomputed_normalised_benchmarks is None:
+        raise ValueError("Must provide either benchmark_texts or precomputed_normalised_benchmarks")
+
     ngrams = get_ngrams(text, n=n)
     if not ngrams:
         return False
 
-    normalised_benchmarks: Optional[Set[str]] = None
-
     for ngram in ngrams:
         if ngram in bloom_filter:
-            if normalised_benchmarks is None:
-                normalised_benchmarks = {_normalise(bt) for bt in benchmark_texts}
+            if precomputed_normalised_benchmarks is None:
+                precomputed_normalised_benchmarks = {_normalise(bt) for bt in (benchmark_texts or [])}
 
-            for nb in normalised_benchmarks:
+            for nb in precomputed_normalised_benchmarks:
                 if f" {ngram} " in f" {nb} ":
                     return True
 

openverifiablellm/utils.py

@@ -1,10 +1,14 @@
+import argparse
 import bz2
 import hashlib
 import json
 import logging
+import os
 import platform
 import re
 import sys
+import time
+import tracemalloc
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -24,6 +28,33 @@
 RE_WHITESPACE = re.compile(r"\s+")
 
 
+# helpers: New helper to compute SHA256 and return raw bytes directly
+def compute_sha256_bytes(
+    *,
+    data: Optional[Union[bytes, bytearray]] = None,
+    file_path: Optional[Union[str, Path]] = None,
+) -> bytes:
+    """
+    Compute SHA256 hash of a file OR raw bytes, returning raw bytes.
+    This avoids the overhead of converting to a hex string and back.
+    """
+    if (data is None) == (file_path is None):
+        raise ValueError("Exactly one of 'data' or 'file_path' must be provided.")
+
+    sha256 = hashlib.sha256()
+
+    if data is not None:
+        sha256.update(data)
+        return sha256.digest()
+
+    path = Path(file_path)
+    with path.open("rb") as f:
+        while chunk := f.read(8192):
+            sha256.update(chunk)
+
+    return sha256.digest()
+
+
 # Merkle Tree Chunk-Level Hashing for Large Files
 def compute_merkle_root(
     file_path: Union[str, Path], chunk_size: int = MERKLE_CHUNK_SIZE_BYTES
@@ -80,8 +111,8 @@ def generate_merkle_proof(
     # Build leaf level
     with path.open("rb") as f:
         while chunk := f.read(chunk_size):
-            leaf_hex = compute_sha256(data=chunk)
-            leaves.append(bytes.fromhex(leaf_hex))
+            leaf_bytes = compute_sha256_bytes(data=chunk)
+            leaves.append(leaf_bytes)
 
     if not leaves:
         raise ValueError("Cannot generate proof for empty file")
@@ -107,8 +138,8 @@ def generate_merkle_proof(
         next_level = []
         for i in range(0, len(leaves), 2):
             combined = leaves[i] + leaves[i + 1]
-            parent_hex = compute_sha256(data=combined)
-            next_level.append(bytes.fromhex(parent_hex))
+            parent_bytes = compute_sha256_bytes(data=combined)
+            next_level.append(parent_bytes)
 
         index //= 2
         leaves = next_level
@@ -121,7 +152,7 @@ def verify_merkle_proof(chunk_bytes: bytes, proof, merkle_root: str) -> bool:
     Verify a Merkle proof for given chunk bytes.
     """
     try:
-        current_hash = bytes.fromhex(compute_sha256(data=chunk_bytes))
+        current_hash = compute_sha256_bytes(data=chunk_bytes)
         expected_root = bytes.fromhex(merkle_root)
     except (TypeError, ValueError):
         return False
@@ -152,8 +183,7 @@ def verify_merkle_proof(chunk_bytes: bytes, proof, merkle_root: str) -> bool:
         else:
             combined = current_hash + sibling
 
-        parent_hex = compute_sha256(data=combined)
-        current_hash = bytes.fromhex(parent_hex)
+        current_hash = compute_sha256_bytes(data=combined)
 
     return current_hash == expected_root
 
@@ -171,9 +201,9 @@ def extract_text_from_xml(
     Process a Wikipedia XML dump (compressed or uncompressed) into cleaned plain text.
 
     Each <page> element is parsed, its revision text is extracted,
-    cleaned using ``clean_wikitext()``, and appended to a single
+    cleaned using clean_wikitext(), and appended to a single
     output text file.
-    When *bloom_filter* and *benchmark_texts* are supplied, each
+    When bloom_filter and benchmark_texts are supplied, each
     cleaned chunk is checked for contamination with evaluation
     benchmarks.  Contaminated chunks are silently skipped.
 
@@ -183,20 +213,9 @@ def extract_text_from_xml(
     Parameters
     ----------
     input_path : str or Path
-        Path to the compressed Wikipedia XML (.bz2) dump file.
-    bloom_filter : rbloom.Bloom, optional
-        Pre-built Bloom filter containing benchmark n-grams.
-    benchmark_texts : list[str], optional
-        Original benchmark text strings for exact verification.
-    n : int, optional
-        N-gram size used for contamination checking (default 13).
-    benchmarks_used : list[str], optional
-        Names of benchmark datasets used (written to manifest).
 
     Output
     ------
-    Creates:
-        data/processed/wiki_clean.txt
     """
     input_path = Path(input_path)
 
@@ -209,9 +228,11 @@ def extract_text_from_xml(
 
     # Lazy import to avoid circular dependency at module level
     if bloom_filter is not None and benchmark_texts is not None:
-        from openverifiablellm.contamination import check_contamination
+        from openverifiablellm.contamination import _normalise, check_contamination
+        precomputed_normalised_benchmarks = {_normalise(bt) for bt in benchmark_texts}
     else:
         check_contamination = None
+        precomputed_normalised_benchmarks = None
 
     # Fixed output path
     project_root = Path.cwd()
@@ -242,8 +263,9 @@ def extract_text_from_xml(
                                 if check_contamination(
                                     cleaned,
                                     bloom_filter,
-                                    benchmark_texts,
+                                    benchmark_texts=None,
                                     n=n,
+                                    precomputed_normalised_benchmarks=precomputed_normalised_benchmarks,
                                 ):
                                     redacted_chunks += 1
                                     elem.clear()
@@ -452,13 +474,89 @@ def clean_wikitext(text: str) -> str:
     return text.strip()
 
 
-if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        print("Usage: python -m openverifiablellm.utils <input_dump> [--no-manifest]")
+def run_benchmark(file_path: str, chunk_size: int = 1024 * 1024):
+    logger.info("--- Starting Benchmark ---")
+
+    if not os.path.exists(file_path):
+        logger.error(f"Error: File not found at {file_path}")
         sys.exit(1)
 
-    logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
-    extract_text_from_xml(
-        sys.argv[1],
-        write_manifest="--no-manifest" not in sys.argv[2:],
+    size_mb = os.path.getsize(file_path) / (1024 * 1024)
+
+    logger.info(f"Benchmarking file: {file_path}")
+    logger.info(f"File size: {size_mb:.2f} MB")
+
+    try:
+        tracemalloc.start()
+
+        # Benchmark compute_merkle_root
+        start_time = time.perf_counter()
+        root_hex = compute_merkle_root(file_path, chunk_size=chunk_size)
+        end_time = time.perf_counter()
+
+        _current_mem, peak_mem = tracemalloc.get_traced_memory()
+
+        root_time = end_time - start_time
+        mins, secs = divmod(root_time, 60)
+        logger.info(f"compute_merkle_root ({size_mb:.2f} MB file): {int(mins)}m {secs:.3f}s")
+        logger.info(f"Peak Memory Usage: {peak_mem / 10**6:.3f} MB")
+        logger.info(f"Merkle Root: {root_hex}")
+
+        tracemalloc.reset_peak()
+
+        # Benchmark generate_merkle_proof
+        start_time = time.perf_counter()
+
+        file_size_bytes = os.path.getsize(file_path)
+        if file_size_bytes == 0:
+            logger.info("Skipping proof benchmark for empty file")
+        else:
+            chunk_count = (file_size_bytes + chunk_size - 1) // chunk_size
+            chunk_index = min(10, chunk_count - 1)
+
+            _ = generate_merkle_proof(file_path, chunk_index=chunk_index, chunk_size=chunk_size)
+            end_time = time.perf_counter()
+
+            _, peak_mem_proof = tracemalloc.get_traced_memory()
+
+            proof_time = end_time - start_time
+            pmins, psecs = divmod(proof_time, 60)
+            logger.info(
+                f"generate_merkle_proof ({size_mb:.2f} MB file, chunk {chunk_index}): {int(pmins)}m {psecs:.3f}s"
+            )
+            logger.info(f"Peak Memory Usage for proof: {peak_mem_proof / 10**6:.3f} MB")
+
+        logger.info("--- Benchmark Complete ---")
+        tracemalloc.stop()
+
+    except Exception:
+        logger.exception("An error occurred during benchmarking")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="OpenVerifiableLLM Preprocessing")
+    parser.add_argument("input_dump", help="Path to the Wikipedia XML dump file")
+    parser.add_argument(
+        "--BENCHMARK_MODE",
+        type=str,
+        choices=["TRUE", "FALSE"],
+        default="FALSE",
+        help="Run in benchmark mode",
     )
+    parser.add_argument(
+        "--chunk_size",
+        type=int,
+        default=MERKLE_CHUNK_SIZE_BYTES,
+        help="Chunk size in bytes for Merkle hashing",
+    )
+    parser.add_argument("--no-manifest", action="store_true", help="Skip manifest generation")
+
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
+
+    if args.BENCHMARK_MODE == "TRUE":
+        run_benchmark(args.input_dump, args.chunk_size)
+    else:
+        extract_text_from_xml(args.input_dump, write_manifest=not args.no_manifest)