Skip to content
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

Handle invalid PyTorch files #30

Merged
merged 4 commits into from
Aug 8, 2023
Merged

Conversation

mmaitre314
Copy link
Contributor

The Hugging Face model repository contains files with a PyTorch .pt extension which are actually PNG files, which trips the picklescan scanner. The code change handles such files as part of the PyTorch magic-number validation and reports them as scan errors instead of unhandled exceptions.

This is a port from mmaitre314/picklescan#19.

@iamfaisalkhan iamfaisalkhan self-requested a review August 7, 2023 15:01
Copy link
Member

@iamfaisalkhan iamfaisalkhan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. Thanks for porting it over! 🎉

bad_pytorch = Modelscan()
bad_pytorch.scan_path(Path(f"{pytorch_file_path}/bad_pytorch.pt"))
assert bad_pytorch.issues.all_issues == []
assert [error.scan_name for error in bad_pytorch.errors] == ["pytorch"] # type: ignore[attr-defined]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not relevant to this PR. scan_name need to be defined as an attribute of the base Error class to avoid the type ignore. I will take care of it in a new PR.

@iamfaisalkhan iamfaisalkhan merged commit 805c35f into protectai:main Aug 8, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants