diff --git a/Legacy/openverifiablellm/utils.py b/Legacy/openverifiablellm/utils.py
index 47cace00..29965b7b 100644
--- a/Legacy/openverifiablellm/utils.py
+++ b/Legacy/openverifiablellm/utils.py
@@ -199,28 +199,65 @@ def _compute_input_identity(input_path: Path) -> str:
return compute_sha256(file_path=input_path)
+def _compute_file_prefix_sha256(file_path: Path, length: int) -> str:
+ """Compute SHA256 of the first `length` bytes of a file."""
+ h = hashlib.sha256()
+ remaining = length
+ with file_path.open("rb") as f:
+ while remaining > 0:
+ chunk = f.read(min(remaining, 65536))
+ if not chunk:
+ break
+ h.update(chunk)
+ remaining -= len(chunk)
+ return h.hexdigest()
+
+
def _load_checkpoint(checkpoint_path: Path, input_path: Path, output_path: Path) -> Dict[str, Any]:
"""Load checkpoint safely and validate resume conditions."""
if not checkpoint_path.exists():
- return {"pages_processed": 0}
+ return {"pages_processed": 0, "file_offset": 0}
try:
with checkpoint_path.open("r", encoding="utf-8") as f:
data = json.load(f)
- pages_processed = data.get("pages_processed")
+ if not isinstance(data, dict):
+ raise ValueError("Checkpoint payload must be a JSON object")
+
+ if "pages_processed" not in data:
+ raise ValueError("Missing pages_processed in checkpoint")
+ pages_processed = data["pages_processed"]
+
+ if "file_offset" not in data:
+ raise ValueError("Missing file_offset in checkpoint")
+ file_offset = data["file_offset"]
+
stored_identity = data.get("input_identity")
+ stored_prefix_hash = data.get("output_prefix_hash")
current_identity = _compute_input_identity(input_path)
- if not isinstance(pages_processed, int) or pages_processed < 0:
+ if not isinstance(pages_processed, int) or isinstance(pages_processed, bool) or pages_processed < 0:
raise ValueError("Invalid pages_processed value")
if stored_identity != current_identity:
raise ValueError("Input file changed since checkpoint")
- if pages_processed > 0 and not output_path.exists():
- raise ValueError("Output file missing; cannot safely resume")
+ if not isinstance(file_offset, int) or isinstance(file_offset, bool) or file_offset < 0:
+ raise ValueError("Invalid file_offset value")
+
+ if pages_processed > 0:
+ if not output_path.exists():
+ raise ValueError("Output file missing; cannot safely resume")
+
+ if output_path.stat().st_size < file_offset:
+ raise ValueError("Output file smaller than checkpoint offset")
+
+ if stored_prefix_hash is not None:
+ actual_prefix_hash = _compute_file_prefix_sha256(output_path, file_offset)
+ if actual_prefix_hash != stored_prefix_hash:
+ raise ValueError("Output prefix digest mismatch; output was modified")
logger.info("Resuming from checkpoint: %d pages already processed", pages_processed)
@@ -228,25 +265,36 @@ def _load_checkpoint(checkpoint_path: Path, input_path: Path, output_path: Path)
except Exception as e:
logger.warning("Checkpoint invalid (%s) — starting fresh.", e)
- return {"pages_processed": 0}
+ return {"pages_processed": 0, "file_offset": 0}
-def _save_checkpoint(checkpoint_path: Path, pages_processed: int, input_identity: str) -> None:
- """Atomically save checkpoint with input identity."""
+def _save_checkpoint(
+ checkpoint_path: Path,
+ pages_processed: int,
+ input_identity: str,
+ file_offset: int = 0,
+ output_prefix_hash: Optional[str] = None,
+) -> None:
+ """Atomically save checkpoint with input identity, file offset, and prefix hash."""
tmp = checkpoint_path.with_suffix(".tmp")
try:
- checkpoint_data = {
+ checkpoint_data: Dict[str, Any] = {
"pages_processed": pages_processed,
"input_identity": input_identity,
+ "file_offset": file_offset,
}
+ if output_prefix_hash is not None:
+ checkpoint_data["output_prefix_hash"] = output_prefix_hash
with tmp.open("w", encoding="utf-8") as f:
json.dump(checkpoint_data, f)
tmp.replace(checkpoint_path)
- logger.debug("Checkpoint saved at %d pages", pages_processed)
+ logger.debug(
+ "Checkpoint saved at %d pages (offset %d)", pages_processed, file_offset
+ )
except Exception as e:
logger.warning("Failed to save checkpoint: %s", e)
@@ -293,9 +341,24 @@ def extract_text_from_xml(input_path, *, write_manifest: bool = False):
# Load checkpoint — tells us how many pages were already written
checkpoint = _load_checkpoint(checkpoint_path, input_path, output_path)
pages_already_done = checkpoint["pages_processed"]
+ file_offset = checkpoint.get("file_offset", 0)
+ input_identity = _compute_input_identity(input_path)
+
+ # If resuming, truncate output file to exact checkpoint offset and append;
+ # otherwise start fresh. Truncating discards any partial writes from crashes
+ # that occurred after the last checkpoint save.
+ if pages_already_done > 0 and output_path.exists():
+ with open(output_path, "r+b") as truncate_f:
+ truncate_f.seek(file_offset)
+ truncate_f.truncate()
+ write_mode = "a"
+ else:
+ write_mode = "w"
+ pages_already_done = 0
+ file_offset = 0
- # If resuming, append to existing output; otherwise start fresh
- write_mode = "a" if pages_already_done > 0 else "w"
+ last_valid_offset = file_offset
+ last_valid_pages = pages_already_done
# Auto-detect file type using magic bytes separation
with open(input_path, "rb") as test_f:
@@ -311,37 +374,83 @@ def extract_text_from_xml(input_path, *, write_manifest: bool = False):
context = ET.iterparse(f, events=("end",))
with open(output_path, write_mode, encoding="utf-8") as out:
- for _, elem in context:
- if elem.tag.endswith("page"):
- pages_seen += 1
+ try:
+ for _, elem in context:
+ if elem.tag.endswith("page"):
+ pages_seen += 1
- # Skip pages already processed in a previous run
- if pages_seen <= pages_already_done:
- elem.clear()
- continue
+ # Skip pages already processed in a previous run
+ if pages_seen <= pages_already_done:
+ elem.clear()
+ continue
- text_elem = elem.find(".//{*}text")
+ text_elem = elem.find(".//{*}text")
- if text_elem is not None and text_elem.text:
- cleaned = clean_wikitext(text_elem.text)
- if cleaned:
- out.write(cleaned + "\n\n")
+ if text_elem is not None and text_elem.text:
+ cleaned = clean_wikitext(text_elem.text)
+ if cleaned:
+ out.write(cleaned + "\n\n")
- pages_written += 1
- elem.clear()
+ pages_written += 1
+ elem.clear()
- # Flush output and save checkpoint periodically
- if pages_written % CHECKPOINT_INTERVAL == 0:
- out.flush()
- _save_checkpoint(checkpoint_path, pages_written, input_path)
- except KeyboardInterrupt:
- _save_checkpoint(checkpoint_path, pages_written, input_path)
- logger.warning("Interrupted by user after %d pages. Run again to resume.", pages_written)
- raise
- except Exception:
- # Save progress before propagating the exception so the next run can resume
- _save_checkpoint(checkpoint_path, pages_written, input_path)
- logger.error("Processing interrupted after %d pages. Run again to resume.", pages_written)
+ # Flush output and save checkpoint periodically
+ if pages_written % CHECKPOINT_INTERVAL == 0:
+ out.flush()
+ last_valid_offset = out.tell()
+ last_valid_pages = pages_written
+ prefix_hash = _compute_file_prefix_sha256(output_path, last_valid_offset)
+ _save_checkpoint(
+ checkpoint_path,
+ pages_written,
+ input_identity,
+ last_valid_offset,
+ prefix_hash,
+ )
+ except KeyboardInterrupt:
+ out.flush()
+ last_valid_offset = out.tell()
+ last_valid_pages = pages_written
+ prefix_hash = _compute_file_prefix_sha256(output_path, last_valid_offset)
+ _save_checkpoint(
+ checkpoint_path,
+ last_valid_pages,
+ input_identity,
+ last_valid_offset,
+ prefix_hash,
+ )
+ logger.warning(
+ "Interrupted by user after %d pages. Run again to resume.", pages_written
+ )
+ raise
+ except Exception:
+ out.flush()
+ last_valid_offset = out.tell()
+ last_valid_pages = pages_written
+ prefix_hash = _compute_file_prefix_sha256(output_path, last_valid_offset)
+ _save_checkpoint(
+ checkpoint_path,
+ last_valid_pages,
+ input_identity,
+ last_valid_offset,
+ prefix_hash,
+ )
+ logger.error(
+ "Processing interrupted after %d pages. Run again to resume.", pages_written
+ )
+ raise
+ except (KeyboardInterrupt, Exception):
+ # If an unhandled exception occurred outside the inner writer (e.g. while opening),
+ # ensure checkpoint retains the last known valid state rather than resetting to 0.
+ if last_valid_pages > 0 and output_path.exists():
+ prefix_hash = _compute_file_prefix_sha256(output_path, last_valid_offset)
+ _save_checkpoint(
+ checkpoint_path,
+ last_valid_pages,
+ input_identity,
+ last_valid_offset,
+ prefix_hash,
+ )
raise
# Processing finished successfully — remove checkpoint so a fresh
diff --git a/Legacy/tests/test_util.py b/Legacy/tests/test_util.py
index a2b2d027..ef5b3c13 100644
--- a/Legacy/tests/test_util.py
+++ b/Legacy/tests/test_util.py
@@ -340,3 +340,171 @@ def test_extract_text_from_xml_malformed_xml(tmp_path, monkeypatch):
with pytest.raises(ET.ParseError):
utils.extract_text_from_xml(input_file)
+
+
+def test_extract_text_from_xml_resume_truncates_and_avoids_duplicates(tmp_path, monkeypatch):
+ xml_content = """
+
+
+ First Page Content
+
+
+ Second Page Content
+
+
+ Third Page Content
+
+
+ Fourth Page Content
+
+
+ """
+ input_file = tmp_path / "simplewiki-pages.xml"
+ input_file.write_text(xml_content, encoding="utf-8")
+ monkeypatch.chdir(tmp_path)
+
+ processed_dir = tmp_path / "data" / "processed"
+ processed_dir.mkdir(parents=True, exist_ok=True)
+ output_file = processed_dir / "wiki_clean.txt"
+ checkpoint_file = processed_dir / "wiki_clean.checkpoint.json"
+
+ # Pre-populate output with pages 1 and 2
+ output_file.write_text("First Page Content\n\nSecond Page Content\n\n", encoding="utf-8")
+ valid_offset = output_file.stat().st_size
+
+ # Simulate an interruption where page 3 was partially written or appended
+ # before the next checkpoint was saved
+ output_file.write_text(
+ "First Page Content\n\nSecond Page Content\n\nThird Page Content (partial/uncommitted)\n\n",
+ encoding="utf-8",
+ )
+
+ # Save checkpoint recorded after 2 pages at valid_offset
+ input_identity = utils._compute_input_identity(input_file)
+ utils._save_checkpoint(checkpoint_file, 2, input_identity, valid_offset)
+
+ # Resume processing
+ utils.extract_text_from_xml(input_file)
+
+ result_text = output_file.read_text(encoding="utf-8")
+ # Verify uncommitted partial text was truncated away and no duplicate pages exist
+ assert "Third Page Content (partial/uncommitted)" not in result_text
+ assert result_text.count("First Page Content") == 1
+ assert result_text.count("Second Page Content") == 1
+ assert result_text.count("Third Page Content") == 1
+ assert result_text.count("Fourth Page Content") == 1
+ assert not checkpoint_file.exists(), "Checkpoint should be removed on successful completion"
+
+
+def test_extract_text_from_xml_resume_invalid_checkpoint_starts_fresh(tmp_path, monkeypatch):
+ xml_content = """
+
+
+ Clean Fresh Run
+
+
+ """
+ input_file = tmp_path / "simplewiki-pages.xml"
+ input_file.write_text(xml_content, encoding="utf-8")
+ monkeypatch.chdir(tmp_path)
+
+ processed_dir = tmp_path / "data" / "processed"
+ processed_dir.mkdir(parents=True, exist_ok=True)
+ output_file = processed_dir / "wiki_clean.txt"
+ checkpoint_file = processed_dir / "wiki_clean.checkpoint.json"
+
+ # Save invalid checkpoint with wrong identity and stale output
+ output_file.write_text("Old Stale Data\n\n", encoding="utf-8")
+ utils._save_checkpoint(checkpoint_file, 5, "wrong_sha256_hash", 100)
+
+ # Should detect invalid checkpoint and start fresh
+ utils.extract_text_from_xml(input_file)
+
+ result_text = output_file.read_text(encoding="utf-8")
+ assert "Old Stale Data" not in result_text
+ assert "Clean Fresh Run" in result_text
+
+
+def test_extract_text_from_xml_checkpoint_boolean_or_prefix_hash_mismatch_starts_fresh(tmp_path, monkeypatch):
+ xml_content = """
+
+
+ Verified Page
+
+
+ """
+ input_file = tmp_path / "simplewiki-pages.xml"
+ input_file.write_text(xml_content, encoding="utf-8")
+ monkeypatch.chdir(tmp_path)
+
+ processed_dir = tmp_path / "data" / "processed"
+ processed_dir.mkdir(parents=True, exist_ok=True)
+ output_file = processed_dir / "wiki_clean.txt"
+ checkpoint_file = processed_dir / "wiki_clean.checkpoint.json"
+
+ # Test boolean offset is rejected and starts fresh
+ input_identity = utils._compute_input_identity(input_file)
+ with checkpoint_file.open("w", encoding="utf-8") as f:
+ json.dump({"pages_processed": 1, "input_identity": input_identity, "file_offset": True}, f)
+
+ output_file.write_text("Corrupted boolean run\n\n", encoding="utf-8")
+ utils.extract_text_from_xml(input_file)
+ assert output_file.read_text(encoding="utf-8") == "Verified Page\n\n"
+
+ # Test prefix hash mismatch is rejected and starts fresh
+ with checkpoint_file.open("w", encoding="utf-8") as f:
+ json.dump(
+ {
+ "pages_processed": 1,
+ "input_identity": input_identity,
+ "file_offset": len("Verified Page\n\n"),
+ "output_prefix_hash": "deadbeef" * 8,
+ },
+ f,
+ )
+ output_file.write_text("Modified file contents\n\n", encoding="utf-8")
+ utils.extract_text_from_xml(input_file)
+ assert output_file.read_text(encoding="utf-8") == "Verified Page\n\n"
+
+
+def test_extract_text_from_xml_exception_saves_valid_offset(tmp_path, monkeypatch):
+ xml_content = """
+
+
+ First Page
+
+
+ Second Page
+
+
+ """
+ input_file = tmp_path / "simplewiki-pages.xml"
+ input_file.write_text(xml_content, encoding="utf-8")
+ monkeypatch.chdir(tmp_path)
+
+ processed_dir = tmp_path / "data" / "processed"
+ processed_dir.mkdir(parents=True, exist_ok=True)
+ checkpoint_file = processed_dir / "wiki_clean.checkpoint.json"
+
+ call_count = 0
+ orig_clean = utils.clean_wikitext
+
+ def failing_clean(text):
+ nonlocal call_count
+ call_count += 1
+ if call_count == 2:
+ raise RuntimeError("Simulated crash on second page")
+ return orig_clean(text)
+
+ monkeypatch.setattr(utils, "clean_wikitext", failing_clean)
+
+ with pytest.raises(RuntimeError, match="Simulated crash"):
+ utils.extract_text_from_xml(input_file)
+
+ assert checkpoint_file.exists()
+ checkpoint_data = json.loads(checkpoint_file.read_text(encoding="utf-8"))
+ assert checkpoint_data["pages_processed"] == 1
+ assert checkpoint_data["file_offset"] > 0
+ assert "output_prefix_hash" in checkpoint_data
+
+