feat: add support for long-context documents - #179
Conversation
|
All contributors have signed the DCO ✍️ ✅ |
eurekayuan
changed the title
eurekayuan
changed the title
|
I have read the DCO document and I hereby sign the DCO. |
Greptile SummaryThis PR adds long-context support across all LLM-heavy stages of the anonymization pipeline by replacing single-call
Confidence Score: 4/5Safe to merge with one gap to address: standalone evaluate() doesn't forward custom window sizing to the final judge. The windowing logic is well-designed throughout — offset rebasing, boundary-span dropping, shrink-on-overflow, and rolling-summary carry-over all work correctly. The one concrete gap is that Anonymizer.evaluate() calls _run_final_judge() without forwarding window_max_render_chars / window_safety_margin_chars, so users who configured a smaller cap to match their LLM's context window won't have it respected during standalone evaluation. Everything in the run() / preview() paths is wired correctly. Files Needing Attention: src/anonymizer/interface/anonymizer.py (evaluate() call site) and src/anonymizer/engine/replace/chunked_replace.py (new_chunk_entities vs chunk_tagged_text on hard-cut boundaries). Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| _clip(summary), | ||
| ) | ||
|
|
||
| stitched = "\n".join(part for part in rewritten_parts if part) |
There was a problem hiding this comment.
iter_boundary_windows, so each tagged[start:end] slice already ends with " ". When the LLM mirrors that structure in its output (natural for paragraph-aware models), every rewritten_chunk also ends with " ", and " ".join(...) then inserts a second newline — producing a blank line between every chunk boundary in the final anonymized document. Joining with "" is sufficient because the delimiter is already part of each chunk.
| stitched = "\n".join(part for part in rewritten_parts if part) | |
| stitched = "".join(part for part in rewritten_parts if part) |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| # Fast path: the full single-call rewrite prompt fits under the cap. | ||
| single_rendered = _render_chunk_prompt(template=params.single_call_prompt_template, chunk_row=row, summary="") | ||
| if len(single_rendered) <= cap: | ||
| logger.debug("rewrite: single-call fast path (rendered=%d chars <= cap=%d)", len(single_rendered), cap) | ||
| text = _rewrite_chunk( | ||
| facade=facade, | ||
| prompt=_compile_template(params.single_call_prompt_template).render(**row), | ||
| system_prompt=params.system_prompt, | ||
| purpose="rewrite-generation", | ||
| ) |
There was a problem hiding this comment.
single_rendered as _render_chunk_prompt(..., summary=""), which prepends the ~270-char continuity preamble, but then the actual LLM call omits that preamble. This means a document whose body-only prompt falls in (cap - 270, cap] chars will be routed into the chunked path unnecessarily. Measure with just the body to match what is actually sent.
| # Fast path: the full single-call rewrite prompt fits under the cap. | |
| single_rendered = _render_chunk_prompt(template=params.single_call_prompt_template, chunk_row=row, summary="") | |
| if len(single_rendered) <= cap: | |
| logger.debug("rewrite: single-call fast path (rendered=%d chars <= cap=%d)", len(single_rendered), cap) | |
| text = _rewrite_chunk( | |
| facade=facade, | |
| prompt=_compile_template(params.single_call_prompt_template).render(**row), | |
| system_prompt=params.system_prompt, | |
| purpose="rewrite-generation", | |
| ) | |
| # Fast path: measure body-only prompt (no continuity preamble) since that is what is sent. | |
| single_rendered = _compile_template(params.single_call_prompt_template).render(**row) | |
| if len(single_rendered) <= cap: | |
| logger.debug("rewrite: single-call fast path (rendered=%d chars <= cap=%d)", len(single_rendered), cap) | |
| text = _rewrite_chunk( | |
| facade=facade, | |
| prompt=single_rendered, | |
| system_prompt=params.system_prompt, | |
| purpose="rewrite-generation", | |
| ) |
| ) | ||
| from anonymizer.engine.ndd.model_loader import resolve_model_alias | ||
| from anonymizer.engine.prompt_utils import substitute_placeholders |
There was a problem hiding this comment.
_compile_template is module-private (underscore-prefixed) in chunked_steps.py. Importing it here creates a hidden coupling: if the function is renamed or inlined, qa_generation.py breaks without any clear contract. Consider exposing it as a public helper in chunked_steps.py or defining a local copy with its own lru_cache in this module.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| row[COL_QUALITY_QA] = _generate(full_rendered, "quality-qa-generation").model_dump() | ||
| return row | ||
|
|
||
| units = json.loads(row.get(COL_MEANING_UNITS_SERIALIZED) or "[]") | ||
| base_len = len(compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: "[]"})) | ||
| batches = _batch_units_by_size(units, base_len, max_render_chars - safety_margin_chars) | ||
| items: list[dict[str, Any]] = [] | ||
| for batch_idx, batch in enumerate(batches): | ||
| rendered = compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: json.dumps(batch, ensure_ascii=False)}) | ||
| out = _generate(rendered, f"quality-qa-generation-batch-{batch_idx}") | ||
| for item in out.items: | ||
| items.append({**item.model_dump(mode="json"), "id": len(items) + 1}) | ||
| row[COL_QUALITY_QA] = QualityQAPairsSchema.model_validate({"items": items}).model_dump() |
There was a problem hiding this comment.
.model_dump() (no mode="json"), while every other windowed generator in this PR consistently uses .model_dump(mode="json"). Without mode="json", Pydantic returns native Python objects rather than JSON-serializable equivalents, which can cause downstream serialization failures. The batched path has the same inconsistency.
| row[COL_QUALITY_QA] = _generate(full_rendered, "quality-qa-generation").model_dump() | |
| return row | |
| units = json.loads(row.get(COL_MEANING_UNITS_SERIALIZED) or "[]") | |
| base_len = len(compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: "[]"})) | |
| batches = _batch_units_by_size(units, base_len, max_render_chars - safety_margin_chars) | |
| items: list[dict[str, Any]] = [] | |
| for batch_idx, batch in enumerate(batches): | |
| rendered = compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: json.dumps(batch, ensure_ascii=False)}) | |
| out = _generate(rendered, f"quality-qa-generation-batch-{batch_idx}") | |
| for item in out.items: | |
| items.append({**item.model_dump(mode="json"), "id": len(items) + 1}) | |
| row[COL_QUALITY_QA] = QualityQAPairsSchema.model_validate({"items": items}).model_dump() | |
| row[COL_QUALITY_QA] = _generate(full_rendered, "quality-qa-generation").model_dump(mode="json") | |
| return row | |
| units = json.loads(row.get(COL_MEANING_UNITS_SERIALIZED) or "[]") | |
| base_len = len(compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: "[]"})) | |
| batches = _batch_units_by_size(units, base_len, max_render_chars - safety_margin_chars) | |
| items: list[dict[str, Any]] = [] | |
| for batch_idx, batch in enumerate(batches): | |
| rendered = compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: json.dumps(batch, ensure_ascii=False)}) | |
| out = _generate(rendered, f"quality-qa-generation-batch-{batch_idx}") | |
| for item in out.items: | |
| items.append({**item.model_dump(mode="json"), "id": len(items) + 1}) | |
| row[COL_QUALITY_QA] = QualityQAPairsSchema.model_validate({"items": items}).model_dump(mode="json") |
| _DEFAULT_MAX_RENDER_CHARS: int = _DetectConfig.model_fields["detection_window_max_render_chars"].default | ||
| _DEFAULT_SAFETY_MARGIN_CHARS: int = _DetectConfig.model_fields["detection_window_safety_margin_chars"].default | ||
|
|
||
|
|
There was a problem hiding this comment.
_first_output calls outputs[0] without checking length. In run_windowed_step with first_only=True, if iter_boundary_windows returns an empty list, outputs is empty and this raises IndexError. The fast path makes this unreachable today, but a defensive guard would make the failure mode explicit.
andreatnvidia
left a comment
andreatnvidia
left a comment
There was a problem hiding this comment.
Thanks for taking this on. This is a substantial first PR, and the overall direction makes sense: split long records into bounded windows, carry forward the state needed for consistency, and keep the replacement map explicit.
I left a few comments on edge cases I think are worth tightening before merge. The main themes are:
- thread the user-supplied window sizing through every windowed stage
- make per-window failures local where possible, instead of dropping the whole record
- validate overlap settings early so a bad config cannot explode into thousands of model calls
- avoid silently accepting empty rewrite chunks as successful output
The tests and docs coverage are in good shape, and I think the feature is close. These changes should make it more reliable on real long documents.
| ), | ||
| *self._qa_wf.columns(selected_models=selected_models), | ||
| *self._rewrite_gen_wf.columns( | ||
| window_max_render_chars=window_max_render_chars, |
There was a problem hiding this comment.
this only threads the user-supplied window cap into rewrite generation. domain classification, sensitivity disposition, QA generation, and final judge still build their window params from module defaults, so a user who lowers Detect.detection_window_max_render_chars still gets ~128k prompts in those stages. Could pass the same kwargs through those columns() calls and _run_final_judge too?
| if params.first_only: | ||
| windows = windows[:1] | ||
| outputs = [] | ||
| for start, end in windows: |
There was a problem hiding this comment.
Claude Code caught this one: once this takes the windowed path, a single transient model error or a chunk that legitimately has no meaning units can drop the whole record. Could wrap each window call, skip/log failed windows, and handle the all-failed case explicitly?
| "prompt scaffolding and tags when sizing augmentation/latent windows." | ||
| ), | ||
| ) | ||
| detection_window_overlap_chars: int = Field( |
There was a problem hiding this comment.
suggestion: can we validate that detection_window_overlap_chars is smaller than the effective window size? Right now overlap == window is accepted and the planners advance one character at a time. My smoke test turned a 20k-char row into 16,001 windows.
| _clip(summary), | ||
| ) | ||
|
|
||
| stitched = "\n".join(part for part in rewritten_parts if part) |
There was a problem hiding this comment.
separate from the newline-stitching comment already here: filtering with if part also hides an empty rewrite chunk. If one window returns {"rewritten_text": ""}, that section disappears with no failed-window count or review signal. Maybe count/flag empty chunks instead of treating them as successful output?
…ate overlap, flag empty rewrite chunks Signed-off-by: eurekayuan <zhuoweny@nvidia.com>
…workflow-column plugins Upstream (#182, #190, 8e6eed5) made the detection workflow exportable for external/distributed executors, requiring every column config to be JSON- serializable. Convert the windowed CustomColumnConfig closures introduced by this branch into plugin column configs (WindowedDetectionConfig, WindowedAugmentationConfig, WindowedLatentConfig) registered via data_designer.plugins entry points, mirroring the chunked-validation plugin. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…repeated content The continuity preamble's 'do NOT repeat earlier content' made the rewriter emit near-empty output for sections resembling earlier ones, collapsing repetitive long documents (~300K chars -> ~3-18K). Instruct each window to be rewritten in full so output length tracks input structure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…> period -> warned hard cut A delimiter occurring only near the start of a window previously produced a degenerate, tiny window (many extra LLM calls on documents with sparse newlines). Boundaries now count only if they keep the window at least half of max_chars; when the primary delimiter has no such occurrence the sentence boundary '.' is tried, and if that also fails the window is hard-cut at max_chars with a warning. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
6639ecb to
051f416
Compare
Stale PR reminderThis PR has had failing checks for 14 days without author activity. Failing checks ( Please push an update or leave a comment if you're still working on this. To stop reminders, add the |
binaryaaron
left a comment
binaryaaron
left a comment
There was a problem hiding this comment.
agent (
review-pr): Long-context correctness and integration review
Reviewed head 051f41609773c22f4a17aeffb388e8c9a8645185.
The long-context implementation has four blocking correctness problems:
- Rewrite windows can split rendered entity tags, causing the affected entity to disappear from the protection disposition.
- Detection drops boundary-touching spans without proving that an overlapping window contains a complete copy. Zero or short overlap can therefore erase PII.
- Partial sensitivity-window success is accepted without validating coverage of every input entity. Those failures do not reach
FailedRecordor human-review state. - Empty rewrite windows are omitted from the stitched document without failing the row or requiring review.
I also found several material defects: the configured prompt cap is not enforced on the final rendered prompt; long-context settings are lost by exported builders and standalone evaluation; final-judge windows become misaligned after length-changing rewrites; and Substitute assigns crossing entities to windows that cannot render their complete spans.
The repeated failure mode is the absence of a shared window contract. Raw, tagged, original, and rewritten offsets are treated as interchangeable, while partial-window outcomes have no typed failure channel. A shared window-plan model, explicit outcome policy, bounded executor, and final-render prompt budget would address most of these defects together.
This branch also needs a rebase before merge. Current main adds CombinedRewriteWorkflow, DataDesigner 0.9, and blocking ty checks. The new plugin column configurations are outside the statically accepted ColumnConfigT, and several generator overrides do not satisfy DataDesigner’s current signatures.
Validation completed:
- Focused tests: 196 passed
- Full test suite: 1,175 passed
- Format, lint, lockfile, SPDX, docs, and diff checks passed
- DCO remains the sole hosted check failure
NDD execution-boundary compliance, no-entity behavior, and bundled-skill compatibility look sound.
|
|
||
| # Chunk on the already-tagged text so tag boundaries stay intact; the delimiter | ||
| # default ("\n") keeps cuts on line breaks. | ||
| windows = iter_boundary_windows(tagged, initial_window, delimiter=params.delimiter) |
There was a problem hiding this comment.
agent (
review-pr): Make rewrite boundaries tag- and span-aware
iter_boundary_windows() can hard-cut at max_chars when no delimiter is available. Here it operates on tagged_text, so the cut can split a tag or entity value. _filter_disposition_to_chunk() then drops the entity’s disposition because neither fragment contains the complete value.
Please plan windows in raw-text/span coordinates and render complete tags afterward, or move intersecting boundaries to safe span edges. Add a delimiter-free test with a protected entity crossing the hard boundary.
| right_artificial = window_end < text_len | ||
| kept: list[EntitySpan] = [] | ||
| for s in spans: | ||
| if left_artificial and s.start_position <= window_start: |
There was a problem hiding this comment.
agent (
review-pr): Do not drop crossing detections unless recovery is proven
drop_boundary_spans() assumes every edge-touching span is recovered intact from an overlapping neighbor, but iter_windows() allows zero overlap and overlap shorter than the entity. In those cases every observation can be dropped before resolve_overlaps().
Please retain boundary provenance and suppress a fragment only after confirming a complete neighboring observation, or enforce a sufficient-overlap contract. Tests should cover overlap 0, overlap shorter than the entity, and entities touching both artificial edges.
| rendered = _compile_template(params.prompt_template).render(**{**row, params.text_column: text[start:end]}) | ||
| try: | ||
| outputs.append(_call(rendered, f"{purpose_prefix}-{start}")) | ||
| except Exception: |
There was a problem hiding this comment.
agent (
review-pr): Reject incomplete sensitivity coverage
A failed window is skipped whenever another window succeeds, but the merged disposition validates only that the result is nonempty. It does not require every input tagged or latent entity to appear. Missing entities are silently excluded from the rewrite instructions, while the surviving row prevents NddAdapter from reporting a FailedRecord.
Please validate the merged keys against the canonical entity set and retry, fail, or explicitly flag incomplete rows.
| # ``if part`` drops empty chunks from the join; ``empty_windows`` records how | ||
| # many so a section silently lost to an empty model response is visible as a | ||
| # review signal rather than passing as successful output. | ||
| stitched = "\n".join(part for part in rewritten_parts if part) |
There was a problem hiding this comment.
agent (
review-pr): Make an empty rewrite window fail or force review
An empty window is omitted from stitched, while _rewrite_empty_windows is only counted and never affects needs_human_review. A document can therefore pass after losing an entire section.
Please fail the row, preserve a marked fallback section, or propagate empty_windows > 0 into human review and the public failure surface.
| # Windowed detection: long documents are tiled into overlapping windows | ||
| # so each detector call stays under the render cap. Serializable plugin | ||
| # config so the exported workflow stays runnable by external executors. | ||
| WindowedDetectionConfig( |
There was a problem hiding this comment.
agent (
review-pr): Resolve the DataDesigner 0.9 plugin typing boundary
With current main’s DataDesigner and blocking ty configuration, the new Windowed*Config objects are rejected by the static ColumnConfigT union. The new model-registry generators also have incompatible generate and agenerate overrides.
Runtime plugin injection does not make that mutation visible to ty. Please introduce a typed project-owned plugin-column boundary or narrow adapter-boundary casts, and make the generator methods satisfy DataDesigner’s current protocol.
| ) | ||
| facade = models[params.alias] | ||
| cap = params.max_render_chars | ||
| initial_window = max(_MIN_WINDOW_CHARS, cap - params.safety_margin_chars) |
There was a problem hiding this comment.
agent (
review-pr): Enforce the cap on the final submitted prompt
The planner sizes raw text before adding tags, replacement state, examples, summaries, continuity text, and response-recipe instructions. The 4,000-character floor can also exceed the configured cap by itself.
Please measure the fully rendered prompt and shrink or replan until the actual input to facade.generate() is within budget. Add captured-prompt assertions for every windowed stage.
| "validator sees per chunk; it is NOT the LLM's context window limit." | ||
| ), | ||
| ) | ||
| detection_window_max_render_chars: int = Field( |
There was a problem hiding this comment.
agent (
review-pr): Give long-context policy a shared owner
These settings are declared on Detect but control replace, rewrite, metadata, and judge stages. Both exported detection builders omit them, and standalone evaluation loses the originating cap.
Please introduce one cross-pipeline long-context policy, serialize it through every builder, and retain enough provenance for evaluate() to apply the same policy.
| n = max(1, math.ceil(total / budget)) | ||
| o_slices = slice_evenly(original, n) | ||
| r_slices = slice_evenly(rewritten, n) | ||
| return [render_judge_prompt(template=template, original=o, rewritten=r) for o, r in zip(o_slices, r_slices)] |
There was a problem hiding this comment.
agent (
review-pr): Align judge windows by content
Independent equal-count character slicing stops corresponding after the rewrite inserts, removes, or reorders content. Sensitive source text and a leaked rewritten copy can land in different judge prompts.
Please align windows using anchors or diff segments, or supply overlapping original context for each rewritten window. Add an early-deletion case with a later leaked identifier.
| summary = "" | ||
| for i, (start, end) in enumerate(windows): | ||
| already = {(r["original"], r["label"]) for r in accumulated} | ||
| chunk_entities = new_chunk_entities(spans, start, end, already) |
There was a problem hiding this comment.
agent (
review-pr): Assign Substitute entities only to complete-span windows
Entity assignment uses the start offset, while tag rendering requires full containment. A boundary-crossing entity is therefore listed in one prompt without its complete tagged value and omitted from the next.
Please move boundaries around spans or carry and reassemble crossing spans before assigning each entity to a prompt.
3 participants
Summary
Several stages embedded the whole document in a single prompt and hit DataDesigner's 512K (
MAX_RENDERED_LEN) render cap, failing outright on long inputs. Every such stage is now windowed: each chunked generator renders its own per-window prompt and calls the model directly, bypassing the cap. Stages keep a single-call fast path when the rendered prompt already fits, so short-document behavior is unchanged.Per-stage windowing
chunked_detection.py, new): Overlapping fixed-size character windows; each window is a raw text slice sent to the detector. Per-window offsets are rebased to global, boundary-touching spans are dropped, and overlaps are resolved (resolve_overlaps).chunked_validation.py): Not a text window — batches candidate entities (≤100 per call), each with a ±500-character excerpt. Calls run in parallel across the validator pool with round-robin + failover. Decisions are merged per row; the row is dropped only if every pool member fails.chunked_augmentation.py): Overlapping character windows over tagged text plus seed JSON. A window dynamically shrinks if its rendered prompt exceeds the cap. Outputs are unioned and deduped by(value, label).chunked_latent.py): Same mechanism as augmentation (rewrite mode only); deduped by(label, value).chunked_replace.py): Abutting newline-aligned windows, no overlap. Each chunk carries the accumulated replacement map and a rolling summary, proposing replacements only for new entities so mappings stay consistent across chunks.chunked_rewrite.py): Abutting newline-aligned windows, no overlap. Runs sequentially, passing a continuity preamble and rolling summary between chunks; rewritten parts are stitched.chunked_final_judge.py, new): Splits original and rewritten text into N positionally-paired slices, scores each, and aggregates per-dimension by minimum. Rubric scales are embedded in the prompt with structured output. Replaces the non-windowedLLMJudgeColumnConfig.Parallel processing
ThreadPoolExecutor; the per-alias rate limit still governs real in-flight calls) and merge afterward.Window sizing
detection_window_max_render_chars(default 128 KiB, clamped ≤ NDD's render cap) is the single knob; it is threaded into detection, augmentation, latent, substitute-map, rewrite, and judge.detection_window_safety_margin_chars(8K) leaves headroom for prompt scaffolding;detection_window_overlap_chars(1K) sets the overlap for the overlapping stages; a 4K floor prevents pathological shrinking.Fault tolerance & failure tracking
trace_dataframe(COL_AUGMENTATION_FAILED_WINDOWS/COL_LATENT_FAILED_WINDOWS); the judge degrades to defaults if all windows fail.Observability
Per-window debug logging across all chunked stages: window ranges/sizes, rendered length vs cap, shrink events, rolling-summary contents, and per-stage entity/replacement/window counts.
Type of Change
Testing
make testpasses locallymake checkpasses locally (format + lint + typecheck + lock-check)Documentation
make docs-buildpasses locally