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

FEAT: add support for local model checkpoints and trust_remote_code in HuggingFaceChatTarget #574

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

KutalVolkan
Copy link
Contributor

Overview

This pull request adds support for loading local model checkpoints in the HuggingFaceChatTarget class through the new model_path parameter. It also introduces the trust_remote_code parameter to handle models that require executing custom code during loading.

These enhancements improve the flexibility and usability of the HuggingFaceChatTarget class by enabling users to:

  • Load models from local directories using the model_path parameter, supporting local model checkpoints.
  • Work with models that include custom code by setting trust_remote_code=True, necessary for models like microsoft/Phi-3-mini-4k-instruct.

Work Completed

  • Updated HuggingFaceChatTarget Class:
    • Added model_path Parameter:
      • Modified the __init__ method to include the optional model_path parameter.
      • Adjusted the load_model_and_tokenizer method to load models from a local directory when model_path is provided.
    • Included trust_remote_code Parameter:
      • Added the trust_remote_code parameter with a default value of False to the __init__ method.
      • Passed trust_remote_code to the from_pretrained methods when loading the tokenizer and model.
    • Ensured Backward Compatibility:
      • Preserved existing functionality for loading models using model_id.
      • Maintained default behavior by setting trust_remote_code=False, so models that do not require it are unaffected.

Related Issue

Issue #572

@KutalVolkan
Copy link
Contributor Author

Hello @romanlutz,

The current demo was just a POC and includes unnecessary elements.
I couldn’t find my old test for HuggingFaceChatTarget to extend the tests—should I write a new one?

Please share your feedback on further enhancements and any other suggestions if needed :)

@romanlutz
Copy link
Contributor

Looks amazing @KutalVolkan ! That was incredibly fast 😅

Tests moved to the targets directory I think. https://github.com/Azure/PyRIT/blob/main/tests/target/test_huggingface_chat_target.py

@KutalVolkan KutalVolkan changed the title [DRAFT] FEAT: add support for local model checkpoints and trust_remote_code in HuggingFaceChatTarget FEAT: add support for local model checkpoints and trust_remote_code in HuggingFaceChatTarget Nov 24, 2024
@KutalVolkan KutalVolkan marked this pull request as ready for review November 24, 2024 10:05
@KutalVolkan
Copy link
Contributor Author

Looks amazing @KutalVolkan ! That was incredibly fast 😅

Tests moved to the targets directory I think. https://github.com/Azure/PyRIT/blob/main/tests/target/test_huggingface_chat_target.py

Thank you, Roman! 😊 It might seem fast, but sometimes it’s just the time zone magic at work. I’m glad we’re keeping things efficient, though! Let me know if there’s anything else I should add.

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.

2 participants