#!/usr/bin/env python3 """Evaluate manifold compression, token budgets, and reconstruction fidelity.""" from __future__ import annotations import argparse import json import math import sys from collections import Counter, defaultdict from dataclasses import dataclass from functools import lru_cache from pathlib import Path from typing import Dict, Iterable, List, Optional, Sequence, Tuple REPO_ROOT = Path(__file__).resolve().parent.parent.parent def _maybe_extend_sys_path() -> None: candidates = [ REPO_ROOT / "score" / "src", REPO_ROOT.parent / "score" / "src", ] for path in candidates: if path.exists(): path_str = str(path) if path_str not in sys.path: sys.path.insert(0, path_str) _maybe_extend_sys_path() from sep_text_manifold import encode, native # type: ignore # noqa: E402 @dataclass class WindowRecord: doc_id: str offset: int signature: str metrics: Dict[str, float] chunk: bytes def _extract_text(record: object, text_key: str) -> str: if isinstance(record, dict): value = record.get(text_key) else: value = None if value is None: raise KeyError(f"Missing text field '{text_key}' in record: {record}") if not isinstance(value, str): value = str(value) return value def _iter_text_directory(root: Path, text_key: str) -> Iterable[Tuple[str, str]]: for path in sorted(root.rglob("*")): if not path.is_file(): continue suffix = path.suffix.lower() if suffix not in {".txt", ".jsonl", ".ndjson", ".json"}: continue relative = path.relative_to(root).as_posix() base = relative[:-len(path.suffix)] if path.suffix else relative base_id = base.replace("/", "__") if suffix == ".txt": yield base_id, path.read_text(encoding="utf-8") elif suffix in {".jsonl", ".ndjson"}: yield from _iter_jsonl_file(path, text_key, doc_prefix=base_id) elif suffix == ".json": yield from _iter_json_file(path, text_key, doc_prefix=base_id) def _iter_jsonl_file(path: Path, text_key: str, doc_prefix: str | None = None) -> Iterable[Tuple[str, str]]: with path.open("r", encoding="utf-8") as handle: for idx, line in enumerate(handle): line = line.strip() if not line: continue record = json.loads(line) text = _extract_text(record, text_key) prefix = doc_prefix or path.stem doc_id = f"{prefix}__{idx:07d}" yield doc_id, text def _iter_json_file(path: Path, text_key: str, doc_prefix: str | None = None) -> Iterable[Tuple[str, str]]: data = json.loads(path.read_text(encoding="utf-8")) prefix = doc_prefix or path.stem if isinstance(data, list): for idx, record in enumerate(data): text = _extract_text(record, text_key) yield f"{prefix}__{idx:07d}", text return if isinstance(data, dict): text = _extract_text(data, text_key) yield prefix, text return raise TypeError(f"Unsupported JSON structure in {path}") def iter_text_documents(root: Path, json_text_key: str = "text") -> Iterable[Tuple[str, str]]: if root.is_dir(): yield from _iter_text_directory(root, json_text_key) return suffix = root.suffix.lower() if suffix in {".jsonl", ".ndjson"}: yield from _iter_jsonl_file(root, json_text_key) return if suffix == ".json": yield from _iter_json_file(root, json_text_key) return # Fallback: treat as plain text file yield root.stem, root.read_text(encoding="utf-8") def sliding_windows(data: bytes, window_bytes: int, stride_bytes: int) -> Iterable[Tuple[int, bytes]]: if not data: return if len(data) <= window_bytes: yield 0, data return for offset in range(0, len(data) - window_bytes + 1, stride_bytes): yield offset, data[offset : offset + window_bytes] tail_start = len(data) - window_bytes if tail_start % stride_bytes != 0: yield tail_start, data[tail_start:] def bits_per_metric(precision: int) -> int: buckets = (10**precision) + 1 return math.ceil(math.log2(buckets)) def signature_storage_bytes(precision: int) -> int: metric_bits = bits_per_metric(precision) metrics_total_bits = metric_bits * 4 # coherence, stability, entropy, hazard λ metrics_bytes = math.ceil(metrics_total_bits / 8) count_bytes = 4 # repetition count return metrics_bytes + count_bytes def build_compressed_representation( text_root: Path, window_bytes: int, stride_bytes: int, precision: int, max_documents: Optional[int] = None, json_text_key: str = "text", document_offset: int = 0, ) -> Tuple[ Dict[str, Dict[str, Dict[str, float]]], Dict[str, List[WindowRecord]], Dict[str, str], Dict[str, int], Dict[str, Dict[str, bytes]], ]: compressed: Dict[str, Dict[str, Dict[str, float]]] = {} doc_windows: Dict[str, List[WindowRecord]] = defaultdict(list) doc_texts: Dict[str, str] = {} doc_sizes: Dict[str, int] = {} prototypes: Dict[str, Dict[str, bytes]] = defaultdict(dict) processed_docs = 0 start_index = max(document_offset, 0) for doc_index, (doc_id, text) in enumerate(iter_text_documents(text_root, json_text_key=json_text_key)): if doc_index < start_index: continue if max_documents is not None and processed_docs >= max_documents: break processed_docs += 1 text_bytes = text.encode("utf-8") doc_texts[doc_id] = text doc_sizes[doc_id] = len(text_bytes) doc_bucket = compressed.setdefault(doc_id, {}) for offset, chunk in sliding_windows(text_bytes, window_bytes, stride_bytes): metrics = encode.encode_window(bytes(chunk)) signature = encode.signature_from_metrics( metrics["coherence"], metrics["stability"], metrics["entropy"], precision=precision, ) entry = doc_bucket.setdefault(signature, {"count": 0.0, "hazard_sum": 0.0}) entry["count"] += 1.0 entry["hazard_sum"] += float(metrics["lambda_hazard"]) chunk_bytes = bytes(chunk) if signature not in prototypes[doc_id]: prototypes[doc_id][signature] = chunk_bytes doc_windows[doc_id].append( WindowRecord( doc_id=doc_id, offset=offset, signature=signature, metrics=metrics, chunk=chunk_bytes, ) ) return compressed, doc_windows, doc_texts, doc_sizes, prototypes def normalise_compressed(compressed: Dict[str, Dict[str, Dict[str, float]]]) -> Dict[str, Dict[str, Dict[str, float]]]: normalised: Dict[str, Dict[str, Dict[str, float]]] = {} for doc_id, signatures in compressed.items(): doc_bucket: Dict[str, Dict[str, float]] = {} for signature, payload in signatures.items(): count = payload["count"] hazard_sum = payload["hazard_sum"] hazard_avg = hazard_sum / count if count else 0.0 doc_bucket[signature] = {"count": count, "hazard": hazard_avg} normalised[doc_id] = doc_bucket return normalised def evaluate_verification( doc_signatures: Dict[str, set], doc_windows: Dict[str, List[WindowRecord]], ) -> Tuple[Dict[str, float], Dict[str, Dict[str, float]], int]: positives = sum(len(windows) for windows in doc_windows.values()) true_positive = positives negatives = 0 false_positive = 0 per_doc_stats: Dict[str, Dict[str, float]] = {} for doc_id, signatures in doc_signatures.items(): doc_pos = len(doc_windows.get(doc_id, [])) doc_true_pos = doc_pos doc_neg = 0 doc_false_pos = 0 for other_id, other_windows in doc_windows.items(): if other_id == doc_id: continue for record in other_windows: doc_neg += 1 negatives += 1 if record.signature in signatures: doc_false_pos += 1 false_positive += 1 precision = doc_true_pos / (doc_true_pos + doc_false_pos) if (doc_true_pos + doc_false_pos) else 1.0 fpr = doc_false_pos / doc_neg if doc_neg else 0.0 recall = 1.0 if doc_pos else 0.0 per_doc_stats[doc_id] = { "positives": doc_pos, "true_positive": doc_true_pos, "negatives": doc_neg, "false_positive": doc_false_pos, "precision": precision, "false_positive_rate": fpr, "recall": recall, "f1": (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0, } precision = true_positive / (true_positive + false_positive) if (true_positive + false_positive) else 1.0 recall = true_positive / positives if positives else 0.0 false_positive_rate = false_positive / negatives if negatives else 0.0 overall = { "positives": positives, "true_positive": true_positive, "negatives": negatives, "false_positive": false_positive, "precision": precision, "recall": recall, "false_positive_rate": false_positive_rate, "f1": (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0, } return overall, per_doc_stats, negatives def reconstruct_document( windows: Sequence[WindowRecord], prototypes: Dict[str, bytes], stride_bytes: int, ) -> bytes: if not windows: return b"" sorted_windows = sorted(windows, key=lambda rec: rec.offset) result = bytearray() for record in sorted_windows: chunk = prototypes.get(record.signature, record.chunk) start = max(record.offset, 0) if len(result) < start: gap = start - len(result) result.extend(chunk[:gap]) overlap = len(result) - start if overlap < 0: overlap = 0 if overlap >= len(chunk): continue result.extend(chunk[overlap:]) return bytes(result) @lru_cache(maxsize=4) def get_tokenizer(name: str, trust_remote_code: bool): from transformers import AutoTokenizer # type: ignore tokenizer = AutoTokenizer.from_pretrained(name, trust_remote_code=trust_remote_code) tokenizer.padding_side = "left" tokenizer.model_max_length = 1_000_000 return tokenizer class TokenizerRuntime: def __init__(self, name: str, trust_remote_code: bool): self.name = name self.trust_remote_code = trust_remote_code self._tokenizer = get_tokenizer(name, trust_remote_code) self._cache: Dict[str, Tuple[int, ...]] = {} def encode(self, text: str) -> Tuple[int, ...]: cached = self._cache.get(text) if cached is not None: return cached tokens = tuple(self._tokenizer.encode(text, add_special_tokens=False)) self._cache[text] = tokens return tokens def count(self, text: str) -> int: return len(self.encode(text)) def levenshtein_distance(a: Sequence[int], b: Sequence[int]) -> int: if len(a) < len(b): a, b = b, a prev = list(range(len(b) + 1)) for i, token_a in enumerate(a, start=1): curr = [i] for j, token_b in enumerate(b, start=1): cost = 0 if token_a == token_b else 1 curr.append( min( curr[-1] + 1, # insertion prev[j] + 1, # deletion prev[j - 1] + cost, ) ) prev = curr return prev[-1] def compute_token_metrics(original: str, reconstructed: str, tokenizer: TokenizerRuntime) -> Dict[str, float]: original_tokens = tokenizer.encode(original) reconstructed_tokens = tokenizer.encode(reconstructed) distance = levenshtein_distance(original_tokens, reconstructed_tokens) original_len = len(original_tokens) reconstructed_len = len(reconstructed_tokens) denom = max(original_len, reconstructed_len, 1) accuracy = 1.0 - (distance / denom) precision = 1.0 - (distance / max(reconstructed_len, 1)) recall = 1.0 - (distance / max(original_len, 1)) f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0 return { "original_tokens": original_len, "reconstructed_tokens": reconstructed_len, "edit_distance": distance, "token_accuracy": accuracy, "token_precision": precision, "token_recall": recall, "token_f1": f1, } def compute_character_metrics(original: str, reconstructed: str) -> Dict[str, float]: original_chars = list(original) reconstructed_chars = list(reconstructed) distance = levenshtein_distance(original_chars, reconstructed_chars) original_len = len(original_chars) reconstructed_len = len(reconstructed_chars) denom = max(original_len, reconstructed_len, 1) accuracy = 1.0 - (distance / denom) precision = 1.0 - (distance / max(reconstructed_len, 1)) recall = 1.0 - (distance / max(original_len, 1)) f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0 normalized_distance = distance / denom return { "original_characters": original_len, "reconstructed_characters": reconstructed_len, "edit_distance": distance, "character_accuracy": accuracy, "character_precision": precision, "character_recall": recall, "character_f1": f1, "normalized_edit_distance": normalized_distance, } def evaluate_manifold( text_root: Path, window_bytes: int, stride_bytes: int, precision: int, tokenizer_name: str = "gpt2", tokenizer_trust_remote_code: bool = False, max_documents: Optional[int] = None, use_native: bool = False, json_text_key: str = "text", document_offset: int = 0, ) -> Dict[str, object]: if use_native: native.set_use_native(True) ( compressed_raw, doc_windows, doc_texts, doc_sizes, prototypes, ) = build_compressed_representation( text_root, window_bytes, stride_bytes, precision, max_documents=max_documents, json_text_key=json_text_key, document_offset=document_offset, ) compressed = normalise_compressed(compressed_raw) doc_signatures = {doc_id: set(bucket.keys()) for doc_id, bucket in compressed.items()} tokenizer_runtime = TokenizerRuntime(tokenizer_name, tokenizer_trust_remote_code) storage_bytes_per_sig = signature_storage_bytes(precision) doc_compressed_size = {doc_id: len(signatures) * storage_bytes_per_sig for doc_id, signatures in doc_signatures.items()} compressed_size = sum(doc_compressed_size.values()) original_size = sum(doc_sizes.values()) compression_ratio = (original_size / compressed_size) if compressed_size else float("inf") signature_doc_counts = Counter() for doc_id, signatures in doc_signatures.items(): for signature in signatures: signature_doc_counts[signature] += 1 shared_signatures = sum(1 for count in signature_doc_counts.values() if count > 1) verification_metrics, per_doc_verification, negatives = evaluate_verification(doc_signatures, doc_windows) per_doc_summary = {} total_text_tokens = 0 total_reconstructed_tokens = 0 total_token_edit_distance = 0 total_stream_tokens = 0 total_unique_tokens = 0 total_characters = 0 total_reconstructed_characters = 0 total_char_edit_distance = 0 for doc_id, windows in doc_windows.items(): original_text = doc_texts[doc_id] reconstructed_bytes = reconstruct_document(windows, prototypes[doc_id], stride_bytes) reconstructed_text = reconstructed_bytes.decode("utf-8", errors="replace") token_metrics = compute_token_metrics(original_text, reconstructed_text, tokenizer_runtime) character_metrics = compute_character_metrics(original_text, reconstructed_text) stream_tokens = len(windows) unique_tokens = len(doc_signatures.get(doc_id, set())) total_text_tokens += token_metrics["original_tokens"] total_reconstructed_tokens += token_metrics["reconstructed_tokens"] total_token_edit_distance += token_metrics["edit_distance"] total_stream_tokens += stream_tokens total_unique_tokens += unique_tokens total_characters += character_metrics["original_characters"] total_reconstructed_characters += character_metrics["reconstructed_characters"] total_char_edit_distance += character_metrics["edit_distance"] per_doc_summary[doc_id] = { "original_size_bytes": doc_sizes[doc_id], "compressed_size_bytes": doc_compressed_size.get(doc_id, 0), "compression_ratio": ( doc_sizes[doc_id] / doc_compressed_size.get(doc_id, 1) if doc_compressed_size.get(doc_id, 0) else float("inf") ), "unique_signatures": unique_tokens, "stream_windows": stream_tokens, "token_metrics": token_metrics, "character_metrics": character_metrics, "normalized_edit_distance": character_metrics["normalized_edit_distance"], "token_compression_unique": ( token_metrics["original_tokens"] / unique_tokens if unique_tokens else float("inf") ), "token_compression_stream": ( token_metrics["original_tokens"] / stream_tokens if stream_tokens else float("inf") ), "verification": per_doc_verification.get(doc_id, {}), } token_compression_unique = ( total_text_tokens / total_unique_tokens if total_unique_tokens else float("inf") ) token_compression_stream = ( total_text_tokens / total_stream_tokens if total_stream_tokens else float("inf") ) token_accuracy = 1.0 - ( total_token_edit_distance / max(total_text_tokens, total_reconstructed_tokens, 1) ) token_precision = 1.0 - (total_token_edit_distance / max(total_reconstructed_tokens, 1)) token_recall = 1.0 - (total_token_edit_distance / max(total_text_tokens, 1)) token_f1 = (2 * token_precision * token_recall / (token_precision + token_recall)) if ( token_precision + token_recall ) else 0.0 normalized_char_edit_distance = ( total_char_edit_distance / max(total_characters, total_reconstructed_characters, 1) ) character_accuracy = 1.0 - normalized_char_edit_distance character_precision = 1.0 - (total_char_edit_distance / max(total_reconstructed_characters, 1)) character_recall = 1.0 - (total_char_edit_distance / max(total_characters, 1)) character_f1 = (2 * character_precision * character_recall / (character_precision + character_recall)) if ( character_precision + character_recall ) else 0.0 try: text_root_rel = str(text_root.relative_to(REPO_ROOT)) except ValueError: text_root_rel = str(text_root) summary: Dict[str, object] = { "text_root": text_root_rel, "json_text_key": json_text_key, "documents": len(doc_windows), "window_bytes": window_bytes, "stride_bytes": stride_bytes, "precision": precision, "tokenizer_name": tokenizer_name, "tokenizer_trust_remote_code": tokenizer_trust_remote_code, "signature_storage_bytes": storage_bytes_per_sig, "original_size_bytes": original_size, "compressed_size_bytes": compressed_size, "compression_ratio": compression_ratio, "unique_signatures": sum(len(signatures) for signatures in doc_signatures.values()), "shared_signatures": shared_signatures, "verification": verification_metrics, "per_document": per_doc_summary, "records": sum(len(windows) for windows in doc_windows.values()), "negatives_evaluated": negatives, "token_metrics": { "text_tokens_total": total_text_tokens, "reconstructed_tokens_total": total_reconstructed_tokens, "token_edit_distance_total": total_token_edit_distance, "manifold_tokens_stream": total_stream_tokens, "manifold_tokens_unique": total_unique_tokens, "token_compression_stream": token_compression_stream, "token_compression_unique": token_compression_unique, "token_accuracy": token_accuracy, "token_precision": token_precision, "token_recall": token_recall, "token_f1": token_f1, }, "character_metrics": { "original_characters_total": total_characters, "reconstructed_characters_total": total_reconstructed_characters, "character_edit_distance_total": total_char_edit_distance, "character_accuracy": character_accuracy, "character_precision": character_precision, "character_recall": character_recall, "character_f1": character_f1, "normalized_edit_distance": normalized_char_edit_distance, }, "use_native": use_native, } return summary def main() -> None: parser = argparse.ArgumentParser(description="Evaluate manifold compression fidelity") parser.add_argument("--text-root", type=Path, required=True, help="Root directory of UTF-8 text files") parser.add_argument("--output", type=Path, required=True, help="Destination JSON summary") parser.add_argument("--window-bytes", type=int, default=256, help="Sliding window size") parser.add_argument("--stride-bytes", type=int, default=192, help="Sliding window stride") parser.add_argument("--precision", type=int, default=2, help="Signature precision (decimal places)") parser.add_argument("--tokenizer", type=str, default="gpt2", help="Tokenizer name or local path") parser.add_argument( "--tokenizer-trust-remote-code", action="store_true", help="Allow remote code when loading the tokenizer (required for some custom tokenizers)", ) parser.add_argument( "--json-text-key", type=str, default="text", help="Field name to read when ingesting JSON/JSONL corpora", ) parser.add_argument("--max-documents", type=int, help="Optional cap on number of documents to process") parser.add_argument("--document-offset", type=int, default=0, help="Skip the first N documents before processing") parser.add_argument("--use-native", action="store_true", help="Prefer the native manifold kernel if available") args = parser.parse_args() text_root = args.text_root.resolve() if not text_root.exists(): raise FileNotFoundError(f"text root not found: {text_root}") summary = evaluate_manifold( text_root=text_root, window_bytes=args.window_bytes, stride_bytes=args.stride_bytes, precision=args.precision, tokenizer_name=args.tokenizer, tokenizer_trust_remote_code=args.tokenizer_trust_remote_code, max_documents=args.max_documents, use_native=args.use_native, json_text_key=args.json_text_key, document_offset=args.document_offset, ) output_path = args.output.resolve() output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") print(json.dumps(summary, indent=2)) if __name__ == "__main__": main()