-
Notifications
You must be signed in to change notification settings - Fork 493
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Automatic framework detection in TasksManager for large models #883
Automatic framework detection in TasksManager for large models #883
Conversation
The documentation is not available anymore as the PR was closed or merged. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Left a few nits, but LGTM besides that!
) | ||
logger.info(f"Local {framework_map[framework]} model found.") | ||
all_files = [ | ||
os.path.relpath(os.path.join(dirpath, file), full_model_path) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I am more in favor to use pathlib here but that's just my tastes:
- os.walk => path.glob
- os.path.relpath => path.relative_to
Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>
…om/fxmarty/optimum into auto-detect-framework-large-models
if any(is_pt_weight_file): | ||
framework = "pt" | ||
elif any(is_tf_weight_file): | ||
framework = "tf" | ||
else: | ||
raise FileNotFoundError( | ||
"Cannot determine framework from given checkpoint location." | ||
f" There should be a {Path(WEIGHTS_NAME).stem}*{Path(WEIGHTS_NAME).suffix} for PyTorch" | ||
f" or {Path(TF2_WEIGHTS_NAME).stem}*{Path(TF2_WEIGHTS_NAME).suffix} for TensorFlow." | ||
) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we should log which framework has been chosen, something like:
logger.info(f"No framework specified so {framework} has been automatically chosen.")
Because here PT will have the priority over TF when there are checkpoints for both frameworks in the repo (which is fine), like BERT for instance: https://huggingface.co/bert-base-uncased/tree/main
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is already logged in the CLI: Framework not specified. Using pt to export to ONNX.
If passing export=True
this is indeed not logged.
Merging as failing tests are unrelated. |
Fixes #867