Add LLM-as-a-judge to eval suite and bump NDD version - #2324
Add LLM-as-a-judge to eval suite and bump NDD version#2324sarahyurick wants to merge 16 commits into
Conversation
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
| if not isinstance(execution, dict): | ||
| raise ValueError("'execution' must be a mapping when provided.") | ||
| mode = execution.get("mode", "single_stage") | ||
| if mode not in {"single_stage", "multi_stage"}: |
There was a problem hiding this comment.
I am not sure if splitting up each judge task into its own NDD stage is worth it, or if it should all stay within the same NDD stage. Both are possible. Not sure it matters though.
There was a problem hiding this comment.
Keeping both for now. I think single_stage is more in-line with existing Curator + NDD patterns, but multi_stage is nice if each individual judges need their own runtime_env and if we want to use resumability between judging stages.
| files_per_partition: int | None, | ||
| ) -> Pipeline: | ||
| """Build one streaming Curator reader → optional language gate → NDD stages → filters → writer pipeline.""" | ||
| # TODO: Add an optional TokenLengthFilter stage before NDD stages so prompts |
There was a problem hiding this comment.
Big TODO. Right now I just do string truncation in the jinja file. But we should decide the best course of action.
Also, since some extractors might return null while others return content, we should decide whether those examples are evaluated or not? Right now I force the empty string and still give it to the judge.
There was a problem hiding this comment.
Opened a Linear issue, since this PR is already fairly large: https://linear.app/nvidia/issue/NMCUR-402/enhance-llm-judge-eval-to-filtermodify-by-token-count-limits.
| # Synthetic Data Generation (SDG) Dependencies | ||
| sdg_cpu = [ | ||
| "data-designer==0.5.5", | ||
| "data-designer==0.9.1", |
There was a problem hiding this comment.
Using this version for my tests. It is only like 2 commits behind NDD main right now.
There was a problem hiding this comment.
cc @ayushdg if you have any opinions about these changes?
| extract_reasoning_content: false | ||
| scores: | ||
| - name: best_extraction | ||
| description: Select the candidate that is most useful as clean document text. |
There was a problem hiding this comment.
I had Codex generate all the Jinja, scoring, etc. prompts. Just a POC for now.
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
ayushdg
left a comment
ayushdg
left a comment
There was a problem hiding this comment.
Minor comments around dependencies. Thanks a lot for the bump!
| "kaldiio; sys_platform == 'never'", | ||
| "levenshtein; sys_platform == 'never'", | ||
| "numba==0.65.0", # Override RAPIDS/legacy caps for the inference image; vLLM 0.22 requires numba 0.65.0 | ||
| "packaging>=25", # Required by data-designer==0.9.1; overrides Lightning's packaging<25 cap from nemo-toolkit[asr] |
There was a problem hiding this comment.
Let's add a followup ticket to bump asr. I think some of these overrides can be avoided.
There was a problem hiding this comment.
Opened a GitHub issue: NVIDIA-NeMo/Speech#16144 and a Linear issue: https://linear.app/nvidia/issue/NMCUR-404/relax-the-lightning240-constraint-for-nemo-speech-asr.
| "xgrammar>=0.1.32", # Override vllm's ==0.1.29 pin to address CVE GHSA-7rgv-gqhr-fxg3 (DoS via multi-layer nesting) | ||
| "sqlfluff>=4.2.0", # Address CVE-2026-46373/46374 (parser DoS); overrides data-designer-engine==0.5.5 sqlfluff<4 cap | ||
| "pandas>=3.0.0, <= 3.0.4" # Override data-designers constraint of <3 | ||
| "sqlfluff>=4.2.0", # Address CVE-2026-46373/46374 (parser DoS); data-designer-engine==0.9.1 supports sqlfluff>=4.1,<5 |
There was a problem hiding this comment.
We can probably get rid of this and do uv lock --upgrade-package sqlfluff to have uv pick up a newer version rather than needing an override now.
| [[package]] | ||
| name = "cryptography" | ||
| version = "50.0.0" | ||
| version = "49.0.0" |
There was a problem hiding this comment.
Can we also override cryptography. Looks like NDD adds a pin but should work with newer versions. It's just following older conventions of keeping an upper pin: https://github.com/NVIDIA-NeMo/DataDesigner/pull/821/changes.
>=50 is recommended for https://github.com/NVIDIA-NeMo/Curator/security/dependabot/217.
We can also open an issue in NDD to remove the upper pin on cryptography
| [[package]] | ||
| name = "gitpython" | ||
| version = "3.1.58" | ||
| version = "3.1.57" |
There was a problem hiding this comment.
Do we know if a specific package is downgrading this?
There was a problem hiding this comment.
Not sure. Looks like the latest is 3.1.60. I can do uv lock --upgrade-package GitPython.
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
VibhuJawa
left a comment
VibhuJawa
left a comment
There was a problem hiding this comment.
Still understanding the workflow. Left an initial review . I wonder if we should do pairwise comparisons as a first pass vs score based judges ?
|
|
||
| ## Input and output | ||
|
|
||
| The runner does not require a fixed text schema. A prompt can reference any fields present in an input JSONL or Parquet row. Keep a stable identifier such as `document_id` or `track_id` when you need to join results to another dataset. |
| } | ||
| ``` | ||
|
|
||
| NDD adds one top-level column for each judge. A judge named `extraction_quality` with a `quality` score produces a result shaped like this: |
There was a problem hiding this comment.
Quick question: Is each judge 1 LLM call or multiple LLM calls ?
| --local-dir /path/to/Qwen3.8-27B | ||
| ``` | ||
|
|
||
| A model that serves fine can still be a bad judge, and the only way to know is to run it. The prompt asks the model in plain language to answer with one of a fixed set of values, like "answer with exactly one of: yes, no, unclear." The pipeline checks that answer against the schema afterward. It doesn't stop the model from answering wrong in the first place — it just drops any row where the answer doesn't match. So a model that doesn't follow instructions well shows up as missing rows, not wrong scores. |
There was a problem hiding this comment.
The thing that helped me identify bad judges if we swap_candidates and results change/hold.
I wonder if we should make that a first class citizen ?
|
|
||
| ## Analyzing results | ||
|
|
||
| Running the same rubric through multiple LLMs turns judge agreement into a signal, not just a sanity check: where the models agree, the record is likely easy and the score can be trusted with less scrutiny; where they disagree, look into why before trusting the rubric or filter at scale. A disagreement can mean the record is genuinely ambiguous or hard to score — evidence for a human-in-the-loop or an `unresolved`-style rubric option — or it can mean the prompt or rubric wording is too vague or underspecified for a model to apply consistently, which calls for tightening the prompt rather than trusting either score. |
There was a problem hiding this comment.
This is also my lived experience
|
|
||
| For the optional FastText language gate, provide `--language`, `--fasttext-langid-model-path`, and optionally `--min-langid-score` and `--language-text-field`. Omitting `--language` skips the stage and does not require FastText. | ||
|
|
||
| Press `Ctrl-C` once to cancel a local Dynamo run and allow normal Ray and inference-server cleanup to finish. If cancellation interrupts cleanup, inspect remaining model-server subprocesses before starting another run. |
There was a problem hiding this comment.
Dont think this is how we/our customers run things. Probably good to clean it up. We should try to keep this llm judge relevant.
| def _load_yaml(path: Path) -> dict[str, object]: | ||
| with path.open(encoding="utf-8") as file: | ||
| config = yaml.safe_load(file) | ||
| if not isinstance(config, dict): | ||
| msg = f"Judge config must contain a mapping: {path}" | ||
| raise TypeError(msg) |
There was a problem hiding this comment.
Should we move this to utils or something ? Feels like a util for us to upstream
|
|
||
|
|
||
| def _place_filters(config: dict[str, object], stages: list[dict[str, object]]) -> list[list[dict[str, object]]]: | ||
| """Place top-level filters after the NDD stage that produces their judge column.""" |
There was a problem hiding this comment.
For my understanding: What is the use of place filters produces their judge column
| def _get_num_workers(config: dict[str, object], *, owner: str) -> int | None: | ||
| """Return an optional fixed Ray worker count for one NDD stage.""" | ||
| num_workers = config.get("num_workers") | ||
| if num_workers is None: | ||
| return None | ||
| if isinstance(num_workers, bool) or not isinstance(num_workers, int) or num_workers <= 0: | ||
| msg = f"{owner} must be a positive integer." | ||
| raise ValueError(msg) | ||
| return num_workers |
There was a problem hiding this comment.
Maybe move to a common util or something ?
| def _build_language_filter_stage( | ||
| *, | ||
| language: str | None, | ||
| model_path: str | None, | ||
| min_score: float, | ||
| text_field: str, | ||
| ) -> ScoreFilter | None: | ||
| """Build an optional FastText language gate without retaining its score column.""" | ||
| if not language: | ||
| return None | ||
| if not model_path: | ||
| msg = "--fasttext-langid-model-path is required when --language is provided." | ||
| raise ValueError(msg) | ||
| if not 0.0 <= min_score <= 1.0: | ||
| msg = "--min-langid-score must be between 0 and 1." | ||
| raise ValueError(msg) | ||
|
|
||
| # FastText is optional, so import it only for jobs that enable this stage. | ||
| from nemo_curator.stages.text.filters.fasttext import FastTextLangId | ||
|
|
||
| return ScoreFilter( | ||
| filter_obj=FastTextLangId( | ||
| model_path=model_path, | ||
| min_langid_score=min_score, | ||
| lang=language, | ||
| ), | ||
| text_field=text_field, | ||
| verbose=True, | ||
| ).with_(name="fasttext_language_filter") | ||
|
|
There was a problem hiding this comment.
I get the motivation for this but i wonder if we want it as a generic feature here ?
| judge_stages = [] | ||
| for stage, filters_after_stage in zip(configured_stages, stage_filters, strict=True): |
There was a problem hiding this comment.
What is the intuition b/w using 1 stage vs multiple stages , do we expect perf to be better b/w them ?
3 participants
Things to consider:
LLMJudgeColumnConfig.multi_modal_context(e.g.,ImageContext) but this still requires usingDocumentBatch(since Curator'sDataDesignerStageis document input/output only).eval/llm_judgedirectory directly.