[Bugfix] Enable essential python files export in llm_ptq examples by adding a customized copy function #351
Conversation
WalkthroughAdds a path-resolving helper and a public Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant U as User
participant HF as hf_ptq.py
participant Utils as example_utils.py
participant FS as Filesystem
U->>HF: run export(source_path, export_path, trust_remote_code)
HF->>HF: perform export (TensorRT-LLM / HF variants)
HF-->>HF: export completes (export_path)
alt trust_remote_code == True
HF->>Utils: copy_custom_model_files(source_path, export_path, True)
Utils->>Utils: _resolve_model_path(...) (local / config / snapshot_download)
Utils->>FS: glob for custom files (configuration_*, modeling_*, tokenization_*, *.json)
FS-->>Utils: matching files
Utils->>FS: copy files (shutil.copy2)
FS-->>Utils: copy results
Utils-->>HF: report copied files / warnings
else trust_remote_code == False
HF-->>U: skip copying custom files
end
HF-->>U: export artifacts ready
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
examples/llm_ptq/hf_ptq.py (2)
19-22: Imports look good; addjsonif parsingconfig.jsonfor auto_map (see below).Needed for the suggested robust copy logic.
-import shutil +import shutil +import json
659-661: Remove duplicate copy call in TRT-LLM branch.You call
copy_custom_model_fileshere and again after both branches, causing redundant I/O and noisy logs.- # Copy custom model files for TensorRT-LLM export as well - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
examples/llm_ptq/hf_ptq.py(4 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: linux
- GitHub Check: wait-checks / wait
- GitHub Check: code-quality
- GitHub Check: build-docs
🔇 Additional comments (1)
examples/llm_ptq/hf_ptq.py (1)
678-680: LGTM to run the copy once post-export.Single call here covers both export paths.
Please sanity‑check with a Phi‑3 repo that uses
trust_remote_code: confirm modeling/config/tokenization files (and any nested package files) appear underexport_path/and that a freshfrom_pretrained(export_path, trust_remote_code=True)succeeds offline.
examples/llm_ptq/hf_ptq.py
Outdated
| def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False): | ||
| """Copy custom model files (configuration_*.py, modeling_*.py, etc.) from source to export directory. | ||
| Args: | ||
| source_path: Path to the original model directory | ||
| export_path: Path to the exported model directory | ||
| trust_remote_code: Whether trust_remote_code was used (only copy files if True) | ||
| """ | ||
| if not trust_remote_code: | ||
| return | ||
|
|
||
| source_dir = Path(source_path) | ||
| export_dir = Path(export_path) | ||
|
|
||
| if not source_dir.exists(): | ||
| print(f"Warning: Source directory {source_path} does not exist") | ||
| return | ||
|
|
||
| if not export_dir.exists(): | ||
| print(f"Warning: Export directory {export_path} does not exist") | ||
| return | ||
|
|
||
| # Common patterns for custom model files that need to be copied | ||
| custom_file_patterns = [ | ||
| "configuration_*.py", | ||
| "modeling_*.py", | ||
| "tokenization_*.py", | ||
| "processing_*.py", | ||
| "image_processing_*.py", | ||
| "feature_extraction_*.py", | ||
| ] | ||
|
|
||
| copied_files = [] | ||
| for pattern in custom_file_patterns: | ||
| for file_path in source_dir.glob(pattern): | ||
| if file_path.is_file(): | ||
| dest_path = export_dir / file_path.name | ||
| try: | ||
| shutil.copy2(file_path, dest_path) | ||
| copied_files.append(file_path.name) | ||
| print(f"Copied custom model file: {file_path.name}") | ||
| except Exception as e: | ||
| print(f"Warning: Failed to copy {file_path.name}: {e}") | ||
|
|
||
| if copied_files: | ||
| print(f"Successfully copied {len(copied_files)} custom model files to {export_path}") | ||
| else: | ||
| print("No custom model files found to copy") | ||
|
|
||
|
|
There was a problem hiding this comment.
Make file copy resilient: handle nested packages via auto_map, preserve structure, and guard same-dir.
Current logic only copies top-level files and may miss packages/modules referenced by HF auto_map (common in custom repos like Phi). It can also double-copy names and doesn’t handle source==export paths.
Apply this refactor:
-def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False):
- """Copy custom model files (configuration_*.py, modeling_*.py, etc.) from source to export directory.
-
- Args:
- source_path: Path to the original model directory
- export_path: Path to the exported model directory
- trust_remote_code: Whether trust_remote_code was used (only copy files if True)
- """
- if not trust_remote_code:
- return
-
- source_dir = Path(source_path)
- export_dir = Path(export_path)
-
- if not source_dir.exists():
- print(f"Warning: Source directory {source_path} does not exist")
- return
-
- if not export_dir.exists():
- print(f"Warning: Export directory {export_path} does not exist")
- return
-
- # Common patterns for custom model files that need to be copied
- custom_file_patterns = [
- "configuration_*.py",
- "modeling_*.py",
- "tokenization_*.py",
- "processing_*.py",
- "image_processing_*.py",
- "feature_extraction_*.py",
- ]
-
- copied_files = []
- for pattern in custom_file_patterns:
- for file_path in source_dir.glob(pattern):
- if file_path.is_file():
- dest_path = export_dir / file_path.name
- try:
- shutil.copy2(file_path, dest_path)
- copied_files.append(file_path.name)
- print(f"Copied custom model file: {file_path.name}")
- except Exception as e:
- print(f"Warning: Failed to copy {file_path.name}: {e}")
-
- if copied_files:
- print(f"Successfully copied {len(copied_files)} custom model files to {export_path}")
- else:
- print("No custom model files found to copy")
+def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False):
+ """Copy custom model Python files from source to export directory.
+ - Supports nested packages referenced by config.json::auto_map
+ - Preserves package structure (subdirectories)
+ - No-op unless trust_remote_code is True
+ """
+ if not trust_remote_code:
+ return
+
+ source_dir = Path(source_path)
+ export_dir = Path(export_path)
+
+ if not source_dir.exists():
+ print(f"Warning: Source directory {source_path} does not exist")
+ return
+ if not export_dir.exists():
+ print(f"Warning: Export directory {export_path} does not exist")
+ return
+ if source_dir.resolve() == export_dir.resolve():
+ print("Info: Source and export directories are identical; skipping custom file copy")
+ return
+
+ # Common patterns for custom model files
+ custom_file_patterns = [
+ "configuration_*.py",
+ "modeling_*.py",
+ "tokenization_*.py",
+ "processing_*.py",
+ "image_processing_*.py",
+ "feature_extraction_*.py",
+ "__init__.py",
+ ]
+
+ candidate_paths: set[Path] = set()
+
+ # 1) Collect files by patterns recursively
+ for pattern in custom_file_patterns:
+ for file_path in source_dir.rglob(pattern):
+ if file_path.is_file():
+ candidate_paths.add(file_path)
+
+ # 2) Collect modules from config.json auto_map if present
+ cfg_path = source_dir / "config.json"
+ if cfg_path.exists():
+ try:
+ with cfg_path.open("r", encoding="utf-8") as f:
+ cfg = json.load(f)
+ auto_map = cfg.get("auto_map", {}) or {}
+
+ def iter_strings(x):
+ if isinstance(x, str):
+ yield x
+ elif isinstance(x, list):
+ for i in x:
+ if isinstance(i, str):
+ yield i
+ elif isinstance(x, dict):
+ for i in x.values():
+ if isinstance(i, str):
+ yield i
+
+ modules = set()
+ for v in auto_map.values():
+ for s in iter_strings(v):
+ s = s.split(":")[0] # optional "module:qualname"
+ mod = s.rsplit(".", 1)[0] # drop class
+ modules.add(mod)
+
+ for mod in modules:
+ rel = Path(mod.replace(".", "/"))
+ mod_file = source_dir / (str(rel) + ".py")
+ pkg_dir = source_dir / rel
+ if mod_file.exists():
+ candidate_paths.add(mod_file)
+ if pkg_dir.is_dir():
+ for p in pkg_dir.rglob("*.py"):
+ candidate_paths.add(p)
+ except Exception as e:
+ print(f"Warning: Failed to parse auto_map from config.json: {e}")
+
+ copied = 0
+ for file_path in sorted(candidate_paths):
+ try:
+ rel = file_path.relative_to(source_dir)
+ dest_path = export_dir / rel
+ dest_path.parent.mkdir(parents=True, exist_ok=True)
+ shutil.copy2(file_path, dest_path)
+ copied += 1
+ print(f"Copied custom model file: {rel}")
+ except Exception as e:
+ print(f"Warning: Failed to copy {file_path}: {e}")
+
+ if copied:
+ print(f"Successfully copied {copied} custom model files to {export_path}")
+ else:
+ print("No custom model files found to copy")Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In examples/llm_ptq/hf_ptq.py around lines 88 to 137, the current
copy_custom_model_files only copies top-level files and can double-copy names or
fail when source==export or when custom code lives in nested packages referenced
via HuggingFace auto_map; update it to detect and preserve directory structure
for both files and packages: skip work when source_path == export_path, traverse
matches for patterns plus any packages referenced in an optional auto_map (or
any subdirectory under source that contains matching files), compute each file's
relative path to source_dir and create the same subdirectory under export_dir
before copying, use shutil.copy2 for files and shutil.copytree for directories
(or a recursive copy that merges into export_dir) while avoiding overwriting
identical existing files (or only overwrite if different), and deduplicate by
tracking already-copied relative paths; ensure parent dirs are created and log
clear warnings on failures.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #351 +/- ##
=======================================
Coverage 73.46% 73.46%
=======================================
Files 172 172
Lines 17640 17640
=======================================
Hits 12959 12959
Misses 4681 4681 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
…thon files which cannot be exported with standard HF methods Signed-off-by: Zhiyu Cheng <[email protected]>
Edwardf0t1
force-pushed
the
zhiyu/fix-phi-3-config-export
branch
from
962656f to
39ccfad
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
examples/llm_ptq/hf_ptq.py (2)
19-22: Consider addingimport jsonif adopting auto_map support.If you implement the recommended auto_map-aware copy below, add a top-level
import json.Example (outside the diff range):
import json
660-662: Remove redundant invocation; call once after both branches.You call copy_custom_model_files again below; drop this earlier call to avoid duplicate work.
- # Copy custom model files for TensorRT-LLM export as well - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
examples/llm_ptq/hf_ptq.py(4 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: linux
- GitHub Check: wait-checks / wait
- GitHub Check: code-quality
- GitHub Check: build-docs
🔇 Additional comments (2)
examples/llm_ptq/hf_ptq.py (2)
679-681: LGTM: single post-export copy site.Keep this as the single copy location after both export paths. With the guard for source==export, this remains safe.
Please confirm that no other code writes to export_path after this call (to avoid overwriting copied files).
88-137: Make copy robust: recurse, preserve package structure, honor auto_map, and guard same-dir.Top-level glob misses nested packages; filenames may collide; copying to the same dir can fail; missing init.py breaks packages. Refactor to rglob, read config.json::auto_map, preserve relative paths, mkdir parents, and skip when source==export.
Apply this diff:
-def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False): - """Copy custom model files (configuration_*.py, modeling_*.py, etc.) from source to export directory. - - Args: - source_path: Path to the original model directory - export_path: Path to the exported model directory - trust_remote_code: Whether trust_remote_code was used (only copy files if True) - """ - if not trust_remote_code: - return - - source_dir = Path(source_path) - export_dir = Path(export_path) - - if not source_dir.exists(): - print(f"Warning: Source directory {source_path} does not exist") - return - - if not export_dir.exists(): - print(f"Warning: Export directory {export_path} does not exist") - return - - # Common patterns for custom model files that need to be copied - custom_file_patterns = [ - "configuration_*.py", - "modeling_*.py", - "tokenization_*.py", - "processing_*.py", - "image_processing_*.py", - "feature_extraction_*.py", - ] - - copied_files = [] - for pattern in custom_file_patterns: - for file_path in source_dir.glob(pattern): - if file_path.is_file(): - dest_path = export_dir / file_path.name - try: - shutil.copy2(file_path, dest_path) - copied_files.append(file_path.name) - print(f"Copied custom model file: {file_path.name}") - except Exception as e: - print(f"Warning: Failed to copy {file_path.name}: {e}") - - if copied_files: - print(f"Successfully copied {len(copied_files)} custom model files to {export_path}") - else: - print("No custom model files found to copy") +def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False): + """Copy custom model Python files from source to export directory. + - Supports nested packages referenced by config.json::auto_map + - Preserves package structure (subdirectories) + - No-op unless trust_remote_code is True + """ + if not trust_remote_code: + return + + source_dir = Path(source_path) + export_dir = Path(export_path) + + if not source_dir.exists(): + print(f"Warning: Source directory {source_path} does not exist") + return + export_dir.mkdir(parents=True, exist_ok=True) + if source_dir.resolve() == export_dir.resolve(): + print("Info: Source and export directories are identical; skipping custom file copy") + return + + custom_file_patterns = [ + "configuration_*.py", + "modeling_*.py", + "tokenization_*.py", + "processing_*.py", + "image_processing_*.py", + "feature_extraction_*.py", + "__init__.py", + ] + + candidate_paths: set[Path] = set() + + # Collect files by patterns recursively + for pattern in custom_file_patterns: + for file_path in source_dir.rglob(pattern): + if file_path.is_file(): + candidate_paths.add(file_path) + + # Collect modules from config.json auto_map if present + cfg_path = source_dir / "config.json" + if cfg_path.exists(): + try: + with cfg_path.open("r", encoding="utf-8") as f: + cfg = json.load(f) + auto_map = cfg.get("auto_map") or {} + + def iter_strings(x): + if isinstance(x, str): + yield x + elif isinstance(x, list): + for i in x: + if isinstance(i, str): + yield i + elif isinstance(x, dict): + for i in x.values(): + if isinstance(i, str): + yield i + + modules = set() + for v in auto_map.values(): + for s in iter_strings(v): + s = s.split(":", 1)[0] # optional "module:qualname" + mod = s.rsplit(".", 1)[0] # drop class name + modules.add(mod) + + for mod in modules: + rel = Path(mod.replace(".", "/")) + mod_file = source_dir / (str(rel) + ".py") + pkg_dir = source_dir / rel + if mod_file.exists(): + candidate_paths.add(mod_file) + if pkg_dir.is_dir(): + for p in pkg_dir.rglob("*.py"): + candidate_paths.add(p) + except Exception as e: + print(f"Warning: Failed to parse auto_map from config.json: {e}") + + copied = 0 + for file_path in sorted(candidate_paths): + try: + rel = file_path.relative_to(source_dir) + dest_path = export_dir / rel + dest_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(file_path, dest_path) + copied += 1 + print(f"Copied custom model file: {rel}") + except Exception as e: + print(f"Warning: Failed to copy {file_path}: {e}") + + if copied: + print(f"Successfully copied {copied} custom model files to {export_path}") + else: + print("No custom model files found to copy")Note: Add
import jsonat the top-level imports.
Signed-off-by: Zhiyu Cheng <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
examples/llm_ptq/example_utils.py (3)
281-291: Guard same-source/destination to avoid SameFileErrorIf source_path == export_path, copy2 will raise SameFileError repeatedly. Add an early no-op.
source_dir = Path(source_path) export_dir = Path(export_path) if not source_dir.exists(): print(f"Warning: Source directory {source_path} does not exist") return if not export_dir.exists(): print(f"Warning: Export directory {export_path} does not exist") return + + # If both paths are the same, skip copying to avoid SameFileError + if source_dir.resolve() == export_dir.resolve(): + print("Info: Source and export directories are identical; skipping custom file copy") + return
292-301: Include init.py in patternsRequired for packages to be importable when custom modules are in subpackages.
custom_file_patterns = [ "configuration_*.py", "modeling_*.py", "tokenization_*.py", "processing_*.py", "image_processing_*.py", "feature_extraction_*.py", + "__init__.py", ]
303-313: Preserve directory structure and recurse to catch nested packagesUse rglob and copy relative paths so nested modules are included and layout is preserved.
- copied_files = [] - for pattern in custom_file_patterns: - for file_path in source_dir.glob(pattern): - if file_path.is_file(): - dest_path = export_dir / file_path.name - try: - shutil.copy2(file_path, dest_path) - copied_files.append(file_path.name) - print(f"Copied custom model file: {file_path.name}") - except Exception as e: - print(f"Warning: Failed to copy {file_path.name}: {e}") + copied_files = [] + for pattern in custom_file_patterns: + for file_path in source_dir.rglob(pattern): + if file_path.is_file(): + rel = file_path.relative_to(source_dir) + dest_path = export_dir / rel + try: + dest_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(file_path, dest_path) + copied_files.append(str(rel)) + print(f"Copied custom model file: {rel}") + except Exception as e: + print(f"Warning: Failed to copy {rel}: {e}")examples/llm_ptq/hf_ptq.py (1)
615-617: Remove duplicate copy invocation in TRT-LLM branchcopy_custom_model_files is called unconditionally below; calling here causes double work and noisy logs.
- # Copy custom model files for TensorRT-LLM export as well - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
examples/llm_ptq/example_utils.py(2 hunks)examples/llm_ptq/hf_ptq.py(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
examples/llm_ptq/hf_ptq.py (1)
examples/llm_ptq/example_utils.py (1)
copy_custom_model_files(270-317)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: linux
- GitHub Check: wait-checks / wait
- GitHub Check: code-quality
- GitHub Check: build-docs
🔇 Additional comments (4)
examples/llm_ptq/example_utils.py (2)
17-20: Imports look goodshutil and Path are appropriate for the new copy utility.
270-317: Consider auto_map-aware copy for completeness (nested module resolution)Parsing config.json::auto_map to copy referenced modules/packages would make this robust for repos like Phi. See prior review suggesting an auto_map-aware implementation.
examples/llm_ptq/hf_ptq.py (2)
26-33: Import looks correctcopy_custom_model_files is added alongside existing utils; consistent with local imports in this module.
634-636: Single post-export copy is sufficientKeeping one invocation after both export branches is the right spot.
Signed-off-by: Zhiyu Cheng <[email protected]>
Signed-off-by: Zhiyu Cheng <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
examples/llm_ptq/example_utils.py (4)
276-345: Avoid executing remote code and fix brittle cache fallback in _resolve_model_pathUsing AutoConfig.from_pretrained can execute remote code when trust_remote_code=True and is slower than needed for path resolution. The TRANSFORMERS_CACHE-based fallback is also brittle and likely incorrect for HF Hub’s cache layout. Prefer snapshot_download when available and otherwise return the input path to fail gracefully.
Apply this diff to simplify and harden resolution:
def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False) -> str: """Resolve a model name or path to a local directory path. If the input is already a local directory, returns it as-is. - If the input is a HuggingFace model ID, attempts to resolve it to the local cache path. + If the input is a HuggingFace model ID, prefer resolving via huggingface_hub without importing model code. @@ - # Try to resolve HuggingFace model ID to local cache path - try: - # First try to load the config to trigger caching - config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code) - - # The config object should have the local path information - # Try different ways to get the cached path - if hasattr(config, "_name_or_path") and os.path.isdir(config._name_or_path): - return config._name_or_path - - # Alternative: use snapshot_download if available - if snapshot_download is not None: - try: - local_path = snapshot_download( - repo_id=model_name_or_path, - allow_patterns=["*.py", "*.json"], # Only download Python files and config - ) - return local_path - except Exception as e: - print(f"Warning: Could not download model files using snapshot_download: {e}") - - # Fallback: try to find in HuggingFace cache - from transformers.utils import TRANSFORMERS_CACHE - - # Look for the model in the cache directory - cache_pattern = os.path.join(TRANSFORMERS_CACHE, "models--*") - cache_dirs = glob.glob(cache_pattern) - - # Convert model name to cache directory format - model_cache_name = model_name_or_path.replace("/", "--") - for cache_dir in cache_dirs: - if model_cache_name in cache_dir: - # Look for the snapshots directory - snapshots_dir = os.path.join(cache_dir, "snapshots") - if os.path.exists(snapshots_dir): - # Get the latest snapshot - snapshot_dirs = [ - d - for d in os.listdir(snapshots_dir) - if os.path.isdir(os.path.join(snapshots_dir, d)) - ] - if snapshot_dirs: - latest_snapshot = max(snapshot_dirs) # Use lexicographically latest - snapshot_path = os.path.join(snapshots_dir, latest_snapshot) - return snapshot_path - - except Exception as e: - print(f"Warning: Could not resolve model path for {model_name_or_path}: {e}") - - # If all else fails, return the original path - # This will cause the copy function to skip with a warning - return model_name_or_path + # Prefer HF Hub snapshot; avoids importing/executing config/model code. + if snapshot_download is not None: + try: + return snapshot_download( + repo_id=model_name_or_path, + allow_patterns=[ + "configuration_*.py", + "modeling_*.py", + "tokenization_*.py", + "processing_*.py", + "image_processing_*.py", + "feature_extraction_*.py", + "*.json", + ], + ) + except Exception as e: + print(f"Warning: snapshot_download failed for {model_name_or_path}: {e}") + # Fall back to the original string; caller will warn/skip if not found locally. + return model_name_or_path
374-377: Create export directory if missing instead of returningThis avoids a no-op when the caller passes a new export path.
- if not export_dir.exists(): - print(f"Warning: Export directory {export_path} does not exist") - return + if not export_dir.exists(): + export_dir.mkdir(parents=True, exist_ok=True)
347-354: *Make patterns extensible without copying all .pyPast feedback suggested copying all .py; keeping a tight default set is good, but allow opt-in expansion.
-def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False): +def copy_custom_model_files( + source_path: str, + export_path: str, + trust_remote_code: bool = False, + extra_patterns: list[str] | None = None, +): @@ - custom_file_patterns = [ + custom_file_patterns = [ "configuration_*.py", "modeling_*.py", "tokenization_*.py", "processing_*.py", "image_processing_*.py", "feature_extraction_*.py", ] + if extra_patterns: + custom_file_patterns.extend(extra_patterns)Also applies to: 379-386
400-404: Return the copied file list to callers (optional)Returning copied_files can help hf_ptq.py log or assert presence in tests.
- if copied_files: - print(f"Successfully copied {len(copied_files)} custom model files to {export_path}") - else: - print("No custom model files found to copy") + if copied_files: + print(f"Successfully copied {len(copied_files)} custom model files to {export_path}") + else: + print("No custom model files found to copy") + return copied_files
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
examples/llm_ptq/example_utils.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: linux
- GitHub Check: wait-checks / wait
- GitHub Check: build-docs
- GitHub Check: code-quality
🔇 Additional comments (2)
examples/llm_ptq/example_utils.py (2)
16-22: LGTM on new filesystem importsglob/shutil/Path additions are appropriate for the new helpers.
30-34: Graceful optional dependency import looks goodDefensive import of snapshot_download is fine and keeps this file usable without huggingface_hub.
| copied_files = [] | ||
| for pattern in custom_file_patterns: | ||
| for file_path in source_dir.glob(pattern): | ||
| if file_path.is_file(): | ||
| dest_path = export_dir / file_path.name | ||
| try: | ||
| shutil.copy2(file_path, dest_path) | ||
| copied_files.append(file_path.name) | ||
| print(f"Copied custom model file: {file_path.name}") | ||
| except Exception as e: | ||
| print(f"Warning: Failed to copy {file_path.name}: {e}") | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Preserve directory structure and recurse; current code may miss files and break imports
- glob() only matches top-level; subpackages like utils/*.py won’t be copied.
- Flattening to export_dir/file.name breaks relative imports if a needed file lives in a subdir.
Apply this diff to recurse and keep relative paths:
- for pattern in custom_file_patterns:
- for file_path in source_dir.glob(pattern):
- if file_path.is_file():
- dest_path = export_dir / file_path.name
- try:
- shutil.copy2(file_path, dest_path)
- copied_files.append(file_path.name)
- print(f"Copied custom model file: {file_path.name}")
- except Exception as e:
- print(f"Warning: Failed to copy {file_path.name}: {e}")
+ for pattern in custom_file_patterns:
+ for file_path in source_dir.rglob(pattern):
+ if file_path.is_file():
+ rel_path = file_path.relative_to(source_dir)
+ dest_path = export_dir / rel_path
+ dest_path.parent.mkdir(parents=True, exist_ok=True)
+ try:
+ shutil.copy2(file_path, dest_path)
+ copied_files.append(str(rel_path))
+ print(f"Copied custom model file: {rel_path}")
+ except Exception as e:
+ print(f"Warning: Failed to copy {rel_path}: {e}")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| copied_files = [] | |
| for pattern in custom_file_patterns: | |
| for file_path in source_dir.glob(pattern): | |
| if file_path.is_file(): | |
| dest_path = export_dir / file_path.name | |
| try: | |
| shutil.copy2(file_path, dest_path) | |
| copied_files.append(file_path.name) | |
| print(f"Copied custom model file: {file_path.name}") | |
| except Exception as e: | |
| print(f"Warning: Failed to copy {file_path.name}: {e}") | |
| copied_files = [] | |
| for pattern in custom_file_patterns: | |
| for file_path in source_dir.rglob(pattern): | |
| if file_path.is_file(): | |
| rel_path = file_path.relative_to(source_dir) | |
| dest_path = export_dir / rel_path | |
| dest_path.parent.mkdir(parents=True, exist_ok=True) | |
| try: | |
| shutil.copy2(file_path, dest_path) | |
| copied_files.append(str(rel_path)) | |
| print(f"Copied custom model file: {rel_path}") | |
| except Exception as e: | |
| print(f"Warning: Failed to copy {rel_path}: {e}") |
🤖 Prompt for AI Agents
In examples/llm_ptq/example_utils.py around lines 388 to 399, the current loop
uses glob() and flattens files into export_dir which misses nested files and
breaks relative imports; change to recurse (use rglob or glob with recursive)
over custom_file_patterns, compute each file's relative path to source_dir,
create the corresponding destination subdirectory under export_dir
(mkdir(parents=True, exist_ok=True)), copy the file to that destination path,
and append the relative path (not just name) to copied_files; also keep the
existing try/except and prints but operate on the preserved relative path.
Signed-off-by: Zhiyu Cheng <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
examples/llm_ptq/example_utils.py (4)
329-337: Choose latest snapshot by mtime, not lexicographic nameUsing max(snapshot_dirs) selects lexicographically, which can point to an older snapshot. Prefer os.path.getmtime.
Apply this diff:
- if snapshot_dirs: - latest_snapshot = max(snapshot_dirs) # Use lexicographically latest - snapshot_path = os.path.join(snapshots_dir, latest_snapshot) - return snapshot_path + if snapshot_dirs: + latest_snapshot = max( + snapshot_dirs, + key=lambda d: os.path.getmtime(os.path.join(snapshots_dir, d)), + ) + return os.path.join(snapshots_dir, latest_snapshot)
315-321: Avoid brittle manual cache scans; prefer HF APIsScanning TRANSFORMERS_CACHE for models--* is fragile and version/env dependent. Consider:
- snapshot_download(..., local_files_only=True, allow_patterns=[".py",".json"]) first, then a networked retry only if needed; or
- using huggingface_hub cache introspection (e.g., scan_cache_dir) to locate a repo snapshot robustly.
This reduces false negatives across environments.
378-381: Optionally create export directory instead of bailingIf the export path is valid but not yet created, you can mkdir(parents=True, exist_ok=True) to proceed.
Apply this diff:
- if not export_dir.exists(): - print(f"Warning: Export directory {export_path} does not exist") - return + if not export_dir.exists(): + export_dir.mkdir(parents=True, exist_ok=True)
382-391: Include package markers to preserve importsWhen copying subpackages, init.py files are needed for classic packages (PEP 420 namespace packages aren’t guaranteed here). Add it to the patterns.
Apply this diff:
custom_file_patterns = [ "configuration_*.py", "modeling_*.py", "tokenization_*.py", "processing_*.py", "image_processing_*.py", "feature_extraction_*.py", + "__init__.py", "*.json", ]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
examples/llm_ptq/example_utils.py(3 hunks)examples/llm_ptq/hf_ptq.py(3 hunks)
✅ Files skipped from review due to trivial changes (1)
- examples/llm_ptq/hf_ptq.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: linux
- GitHub Check: wait-checks / wait
- GitHub Check: build-docs
- GitHub Check: code-quality
🔇 Additional comments (5)
examples/llm_ptq/example_utils.py (5)
16-22: LGTM: required stdlib importsImports for glob/shutil/Path are appropriate for the new utilities.
30-34: LGTM: optional huggingface_hub dependencyGraceful ImportError handling for snapshot_download is good.
359-361: LGTM: gate copying on trust_remote_codeEarly return keeps behavior scoped to the intended use-case.
394-407: Preserve directory structure and recurse; current code may miss files and break importsglob() is non‑recursive and flattening into export_dir/file.name can:
- Miss files in subpackages (e.g., utils/*.py).
- Break relative/package imports and cause runtime failures.
Also, copy2 overwrites silently; avoid clobbering existing exports.Apply this diff:
- for pattern in custom_file_patterns: - for file_path in source_dir.glob(pattern): - if file_path.is_file(): - # Skip config.json and model.safetensors.index.json as they're handled separately - if file_path.name in ["config.json", "model.safetensors.index.json"]: - continue - dest_path = export_dir / file_path.name - try: - shutil.copy2(file_path, dest_path) - copied_files.append(file_path.name) - print(f"Copied custom model file: {file_path.name}") - except Exception as e: - print(f"Warning: Failed to copy {file_path.name}: {e}") + for pattern in custom_file_patterns: + for file_path in source_dir.rglob(pattern): + if file_path.is_file(): + # Skip config.json and model.safetensors.index.json as they're handled separately + if file_path.name in ["config.json", "model.safetensors.index.json"]: + continue + rel_path = file_path.relative_to(source_dir) + dest_path = export_dir / rel_path + dest_path.parent.mkdir(parents=True, exist_ok=True) + try: + if dest_path.exists(): + print(f"Skipping existing file (no overwrite): {rel_path}") + continue + shutil.copy2(file_path, dest_path) + copied_files.append(str(rel_path)) + print(f"Copied custom model file: {rel_path}") + except Exception as e: + print(f"Warning: Failed to copy {rel_path}: {e}")
347-358: Confirm helper invocation in hf_ptq.py Invocation confirmed at lines 616 and 635 usingargs.pyt_ckpt_path,export_path, andargs.trust_remote_code; no changes needed.
…adding a customized copy function (#351) Signed-off-by: Zhiyu Cheng <[email protected]>
…adding a customized copy function (#351) Signed-off-by: Zhiyu Cheng <[email protected]> Signed-off-by: Ye Yu <[email protected]>
What does this PR do?
Type of change: Bug fix
Overview: Add a
copy_custom_model_filesfunction to copy essential python files which cannot be exported with standard HF methods.Resolve https://nvbugspro.nvidia.com/bug/5528021
Usage
Testing
Before your PR is "Ready for review"
Additional Information
Summary by CodeRabbit
New Features
Chores