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

Add support for prompt augmentation #766

Merged
merged 6 commits into from
Mar 11, 2024

Conversation

k-ivey
Copy link
Collaborator

@k-ivey k-ivey commented Dec 14, 2023

What does this PR do?

Summary

This PR adds support for augmenting prompts in TextAttack as well as using these augmented prompts to run LLMs. Note that the design is slightly different than previous offline discussions. Notably, this change reuses the existing Augmenter class to leverage existing recipes and code. An example (copied from a new test) is below:

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

from textattack.augmentation.recipes import CheckListAugmenter
from textattack.constraints.pre_transformation import UnmodifiableIndices
from textattack.llms import HuggingFaceLLMWrapper
from textattack.prompt_augmentation import PromptAugmentationPipeline

model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
model_wrapper = HuggingFaceLLMWrapper(model, tokenizer)

augmenter = CheckListAugmenter()

pipeline = PromptAugmentationPipeline(augmenter, model_wrapper)

prompt = "As a sentiment classifier, determine whether the following text is 'positive' or 'negative'. Please classify: Poor Ben Bratt couldn't find stardom if MapQuest emailed him point-to-point driving directions."
prompt_constraints = [UnmodifiableIndices([2, 3, 10, 12, 14])]

output = pipeline(prompt, prompt_constraints)

Additions

  • Added llms module with a wrapper for HuggingFace and OpenAI
  • Added prompt_augmentation module with a pipeline to augment a prompt and run the augmented prompts on a LLM
  • Added two new constraints UnmodifiableIndices and UnmodifiablePhrases to give users the options to restrict certain parts of a prompt from being augmented

Changes

  • Updated AttackedText to fix a bug in convert_from_original_idxs as the type argument in ininstance must be a tuple if multiple types are provided

To-Do

  • Add CLI support
  • Potentially add additional documentation per feedback

Checklist

  • [ X ] The title of your pull request should be a summary of its contribution.
  • [ X ] Please write detailed description of what parts have been newly added and what parts have been modified. Please also explain why certain changes were made.
  • [ ] If your pull request addresses an issue, please mention the issue number in the pull request description to make sure they are linked (and people consulting the issue know you are working on it)
  • [ X ] To indicate a work in progress please mark it as a draft on Github.
  • [ X ] Make sure existing tests pass.
  • [ X ] Add relevant tests. No quality testing = no merge.
  • [ X ] All public methods must have informative docstrings that work nicely with sphinx. For new modules/files, please add/modify the appropriate .rst file in TextAttack/docs/apidoc.'

@qiyanjun qiyanjun marked this pull request as ready for review December 26, 2023 21:19
@qiyanjun
Copy link
Member

@k-ivey is this ready? not a DRAFT stage right?

@qiyanjun qiyanjun self-requested a review March 10, 2024 19:38
@qiyanjun qiyanjun changed the title [DRAFT] Add support for prompt augmentation Add support for prompt augmentation Mar 11, 2024
@qiyanjun
Copy link
Member

Local pytest all passed. so merge!
=================== 72 passed, 12 warnings in 2450.56s (0:40:50) =====================================================

@qiyanjun qiyanjun merged commit a8dfcb1 into QData:master Mar 11, 2024
6 checks passed
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