[#5247][fix] auto-detect local cnn_dailymail dataset by directory layout

[guan404ming]

Summary by CodeRabbit

• New Features

  • Enhanced CNN/DailyMail dataset recognition to support both named datasets and local repository structures.

• Tests

  • Added comprehensive unit tests validating dataset structure detection across multiple filesystem scenarios.

Description

Fixes #5247

get_calib_dataloader detects cnn_dailymail by substring-matching the calibration path. If the local directory name doesn't contain "cnn_dailymail", detection fails and load_dataset raises ValueError: Config name is missing.

Add _is_cnn_dailymail_local_repo as a directory-layout fallback (version subdirs or cnn_dailymail.py builder). Applied to both HF and NeMo calibration loaders.

Test Coverage

tests/unittest/others/test_quantize_calib_dataset.py covers each detection branch plus a negative case.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

Added a helper function _is_cnn_dailymail_local_repo() that detects local CNN/DailyMail dataset directories by checking for version subdirectories or a builder script, extending dataset recognition logic in get_calib_dataloader() and get_nemo_calib_dataloader() beyond substring matching.

Changes

Dataset Detection Enhancement

Layer / File(s)	Summary
Detection Logic tensorrt_llm/quantization/quantize_by_modelopt.py	New _is_cnn_dailymail_local_repo() helper detects local CNN/DailyMail repos by checking for version subdirectories (e.g., 1.0.0) or a cnn_dailymail.py script.
Integration tensorrt_llm/quantization/quantize_by_modelopt.py	get_calib_dataloader() and get_nemo_calib_dataloader() now use the helper alongside substring matching to recognize CNN/DailyMail datasets in any path.
Tests tests/unittest/others/test_quantize_calib_dataset.py	New test test_is_cnn_dailymail_local_repo() validates the helper across five scenarios: non-existent path, empty directory, version subdirectories, builder script presence, and unrelated content.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly and specifically describes the main change: improving auto-detection of local CNN/DailyMail datasets via directory layout inspection.
Description check	✅ Passed	The description provides clear context about the issue, explains the solution with specifics about the fallback mechanism, lists test coverage, and confirms PR checklist completion.
Linked Issues check	✅ Passed	The changes fully address issue #5247 by implementing directory-layout detection to recognize CNN/DailyMail datasets without substring matching, with comprehensive test coverage.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to detecting local CNN/DailyMail datasets: a helper function, updates to calibration loaders, copyright refresh, and targeted test coverage.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Review rate limit: 9/10 reviews remaining, refill in 6 minutes.}

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

tests/unittest/others/test_quantize_calib_dataset.py (2)

21-47: ⚡ Quick win

Add loader-level tests so integration regressions are caught

This test validates _is_cnn_dailymail_local_repo() directly, but it won’t fail if either get_calib_dataloader() or get_nemo_calib_dataloader() stops using that helper. Please add one focused test per loader for a local dir without "cnn_dailymail" in its name but with a valid layout.
As per coding guidelines, “Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix.”

Suggested test additions (example)

+def test_get_calib_dataloader_uses_local_cnn_layout(monkeypatch, tmp_path):
+    # local dir name intentionally doesn't contain cnn_dailymail
+    (tmp_path / "3.0.0").mkdir()
+
+    from tensorrt_llm.quantization import quantize_by_modelopt as qbm
+
+    calls = {}
+    def fake_load_dataset(path, name=None, split=None, trust_remote_code=None):
+        calls["path"] = path
+        calls["name"] = name
+        return {"article": ["hello world"]}
+
+    class _Tok:
+        def batch_encode_plus(self, texts, return_tensors, padding, truncation, max_length):
+            import torch
+            return {
+                "input_ids": torch.tensor([[1, 2, 3]]),
+                "attention_mask": torch.tensor([[1, 1, 1]]),
+            }
+
+    monkeypatch.setattr(qbm, "load_dataset", fake_load_dataset)
+    dl = qbm.get_calib_dataloader(
+        dataset_name_or_dir=str(tmp_path),
+        tokenizer=_Tok(),
+        calib_size=1,
+        batch_size=1,
+        block_size=8,
+    )
+    _ = next(iter(dl))
+    assert calls["name"] == "3.0.0"
+
+
+def test_get_nemo_calib_dataloader_uses_local_cnn_layout(monkeypatch, tmp_path):
+    (tmp_path / "cnn_dailymail.py").write_text("# stub")
+    from tensorrt_llm.quantization import quantize_by_modelopt as qbm
+
+    def fake_load_dataset(path, name=None, split=None, trust_remote_code=None):
+        return {"article": ["abc"]}
+
+    monkeypatch.setattr(qbm, "load_dataset", fake_load_dataset)
+    batches = list(
+        qbm.get_nemo_calib_dataloader(
+            dataset_name_or_dir=str(tmp_path),
+            batch_size=1,
+            calib_size=1,
+            max_sequence_length=16,
+        )
+    )
+    assert batches == [["abc"]]

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/others/test_quantize_calib_dataset.py` around lines 21 - 47,
The current unit test only verifies _is_cnn_dailymail_local_repo(), but we need
integration tests to ensure get_calib_dataloader() and
get_nemo_calib_dataloader() actually use that helper for local dirs; add one
test per loader that creates a temporary directory NOT named "cnn_dailymail" but
with a valid cnn_dailymail layout (e.g., a version subdirectory or a
cnn_dailymail.py file), then call get_calib_dataloader(...) and
get_nemo_calib_dataloader(...) with that path and assert they return a
dataloader/loader instance (or expected behavior) to catch regressions where the
helper stops being used; reference _is_cnn_dailymail_local_repo,
get_calib_dataloader, and get_nemo_calib_dataloader when locating code to
modify.

---

21-47: QA test-list updates are not needed for this PR

Since this change is unit-test scoped (tests/unittest/...) and does not add/alter integration defs, no update to tests/integration/test_lists/qa/* is required.

As per coding guidelines, QA list updates are generally for scheduled integration/perf/Triton runs, not unit-test-only changes.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/others/test_quantize_calib_dataset.py` around lines 21 - 47,
The change is unit-test scoped
(tests/unittest/others/test_quantize_calib_dataset.py) for the function
test_is_cnn_dailymail_local_repo and does not require any updates to integration
QA lists; do not modify or add entries under tests/integration/test_lists/qa/*
and ensure the test remains under the unittest directory (no changes to
integration/perf/Triton scheduling config).

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/quantization/quantize_by_modelopt.py`:
- Line 1: The file tensorrt_llm/quantization/quantize_by_modelopt.py is marked
executable but lacks a shebang (Ruff EXE002); for a library module remove the
executable permission instead of adding a shebang—update the file mode (e.g.,
git update-index --chmod=-x tensorrt_llm/quantization/quantize_by_modelopt.py or
chmod -x locally and commit) so the file is no longer executable.

---

Nitpick comments:
In `@tests/unittest/others/test_quantize_calib_dataset.py`:
- Around line 21-47: The current unit test only verifies
_is_cnn_dailymail_local_repo(), but we need integration tests to ensure
get_calib_dataloader() and get_nemo_calib_dataloader() actually use that helper
for local dirs; add one test per loader that creates a temporary directory NOT
named "cnn_dailymail" but with a valid cnn_dailymail layout (e.g., a version
subdirectory or a cnn_dailymail.py file), then call get_calib_dataloader(...)
and get_nemo_calib_dataloader(...) with that path and assert they return a
dataloader/loader instance (or expected behavior) to catch regressions where the
helper stops being used; reference _is_cnn_dailymail_local_repo,
get_calib_dataloader, and get_nemo_calib_dataloader when locating code to
modify.
- Around line 21-47: The change is unit-test scoped
(tests/unittest/others/test_quantize_calib_dataset.py) for the function
test_is_cnn_dailymail_local_repo and does not require any updates to integration
QA lists; do not modify or add entries under tests/integration/test_lists/qa/*
and ensure the test remains under the unittest directory (no changes to
integration/perf/Triton scheduling config).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3501cfed-4c1d-4b10-904e-6f00eeb1e269

📥 Commits

Reviewing files that changed from the base of the PR and between fdf7bf4 and 63a757b.

📒 Files selected for processing (2)

• tensorrt_llm/quantization/quantize_by_modelopt.py
• tests/unittest/others/test_quantize_calib_dataset.py

[karljang]

@longlee0622 , could you please review this PR?
My Claude suggests you as the appropriate person to review it, but please feel free to reassign someone else on the team if you prefer.

[longlee0622]

/bot run

[longlee0622]

/bot run

[tensorrt-cicd]

PR_Github #51794 [ run ] triggered by Bot. Commit: dc2428c Link to invocation

[tensorrt-cicd]

PR_Github #51793 [ run ] triggered by Bot. Commit: dc2428c Link to invocation

[tensorrt-cicd]

PR_Github #51793 [ run ] completed with state ABORTED. Commit: dc2428c

Link to invocation

[tensorrt-cicd]

PR_Github #51794 [ run ] completed with state ABORTED. Commit: dc2428c
/LLM/main/L0_MergeRequest_PR pipeline #41157 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[longlee0622]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51813 [ run ] triggered by Bot. Commit: dc2428c Link to invocation

[tensorrt-cicd]

PR_Github #51813 [ run ] completed with state SUCCESS. Commit: dc2428c
/LLM/main/L0_MergeRequest_PR pipeline #41177 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[guan404ming]

Thanks @longlee0622!