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.

Sign up for GitHub

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

Feature: use hf chat support #1047

Merged
merged 6 commits into from
Dec 19, 2024
Merged

Feature: use hf chat support #1047

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
87ab1ca
detect chat support from tokenizer
jmartin-tech Dec 11, 2024
9449258
test a chat enabled hf model
jmartin-tech Dec 16, 2024
9ec5794
always deprefix when using chat template
jmartin-tech Dec 17, 2024
99c760a
use pipeline explicit truncation, use tokenizer decode in Model
jmartin-tech Dec 17, 2024
900fc8f
allow user provided `use_chat` override
jmartin-tech Dec 19, 2024
86de116
ensure output format on chat enabled model tests
jmartin-tech Dec 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 59 additions & 12 deletions garak/generators/huggingface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,9 @@ def _load_client(self):
set_seed(_config.run.seed)

pipeline_kwargs = self._gather_hf_params(hf_constructor=pipeline)
pipeline_kwargs["truncation"] = (
Copy link
Collaborator

Choose a reason for hiding this comment

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

Sometimes I wish this were just the default for HF but c'est la vie. We may want to consider whether truncation = True requires max_len to be set -- I've had HF yell at me for not specifying both but it may be a corner case that I encountered.

Copy link
Collaborator

Choose a reason for hiding this comment

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

HF is kinda yell-y, and over the course of garak dev, HF has varied what it yells about. It can also be the case that some models require certain param combos to operate, while others will find the same param combo utterly intolerable. This has led to a style where one tries to do the right thing in garak, and tries to listen to HF warnings a little less.

True # this is forced to maintain existing pipeline expectations
)
self.generator = pipeline("text-generation", **pipeline_kwargs)
if self.generator.tokenizer is None:
# account for possible model without a stored tokenizer
Expand All @@ -87,6 +90,11 @@ def _load_client(self):
self.generator.tokenizer = AutoTokenizer.from_pretrained(
pipeline_kwargs["model"]
)
if not hasattr(self, "use_chat"):
self.use_chat = (
hasattr(self.generator.tokenizer, "chat_template")
and self.generator.tokenizer.chat_template is not None
)
if not hasattr(self, "deprefix_prompt"):
self.deprefix_prompt = self.name in models_to_deprefix
if _config.loaded:
Expand All @@ -98,6 +106,9 @@ def _load_client(self):
def _clear_client(self):
self.generator = None

def _format_chat_prompt(self, prompt: str) -> List[dict]:
return [{"role": "user", "content": prompt}]

def _call_model(
self, prompt: str, generations_this_call: int = 1
) -> List[Union[str, None]]:
Expand All @@ -106,13 +117,16 @@ def _call_model(
warnings.simplefilter("ignore", category=UserWarning)
try:
with torch.no_grad():
# workaround for pipeline to truncate the input
encoded_prompt = self.generator.tokenizer(prompt, truncation=True)
truncated_prompt = self.generator.tokenizer.decode(
encoded_prompt["input_ids"], skip_special_tokens=True
)
# according to docs https://huggingface.co/docs/transformers/main/en/chat_templating
# chat template should be automatically utilized if the pipeline tokenizer has support
# and a properly formatted list[dict] is supplied
if self.use_chat:
formatted_prompt = self._format_chat_prompt(prompt)
else:
formatted_prompt = prompt
Comment on lines +123 to +126
Copy link
Collaborator

Choose a reason for hiding this comment

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

Note for future Erick and Jeffrey: we will want to merge ConversationalPipeline into this at some point given HF's deprecation of the "conversational" pipeline and Conversation object. At that point, we'll also want to validate that prompt is a string.


raw_output = self.generator(
truncated_prompt,
formatted_prompt,
pad_token_id=self.generator.tokenizer.eos_token_id,
max_new_tokens=self.max_tokens,
num_return_sequences=generations_this_call,
Expand All @@ -127,10 +141,15 @@ def _call_model(
i["generated_text"] for i in raw_output
] # generator returns 10 outputs by default in __init__

if self.use_chat:
text_outputs = [_o[-1]["content"].strip() for _o in outputs]
else:
text_outputs = outputs

Comment on lines +144 to +148
Copy link
Collaborator

@leondz leondz Dec 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

Comment not for this PR: there are multiple ways of representing conversations in garak, and attempt should be canonical for conversation history; I'm starting to consider patterns where attempt holds the conversation history still but where this can be read & written using other interfaces, like an OpenAI API messages dict list, or this HF style format. But mayeb the work is so simple that this current pattern works fine.

if not self.deprefix_prompt:
return outputs
return text_outputs
else:
return [re.sub("^" + re.escape(prompt), "", _o) for _o in outputs]
return [re.sub("^" + re.escape(prompt), "", _o) for _o in text_outputs]


class OptimumPipeline(Pipeline, HFCompatible):
Expand Down Expand Up @@ -468,6 +487,13 @@ def _load_client(self):
self.name, padding_side="left"
)

if not hasattr(self, "use_chat"):
# test tokenizer for `apply_chat_template` support
self.use_chat = (
hasattr(self.tokenizer, "chat_template")
and self.tokenizer.chat_template is not None
)

self.generation_config = transformers.GenerationConfig.from_pretrained(
self.name
)
Expand All @@ -492,14 +518,27 @@ def _call_model(
if self.top_k is not None:
self.generation_config.top_k = self.top_k

text_output = []
raw_text_output = []
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=UserWarning)
with torch.no_grad():
if self.use_chat:
formatted_prompt = self.tokenizer.apply_chat_template(
self._format_chat_prompt(prompt),
tokenize=False,
add_generation_prompt=True,
)
else:
formatted_prompt = prompt

inputs = self.tokenizer(
prompt, truncation=True, return_tensors="pt"
formatted_prompt, truncation=True, return_tensors="pt"
).to(self.device)

prefix_prompt = self.tokenizer.decode(
inputs["input_ids"][0], skip_special_tokens=True
)

try:
outputs = self.model.generate(
**inputs, generation_config=self.generation_config
Expand All @@ -512,14 +551,22 @@ def _call_model(
return returnval
else:
raise e
text_output = self.tokenizer.batch_decode(
raw_text_output = self.tokenizer.batch_decode(
outputs, skip_special_tokens=True, device=self.device
)

if self.use_chat:
text_output = [
re.sub("^" + re.escape(prefix_prompt), "", i).strip()
for i in raw_text_output
]
else:
text_output = raw_text_output
Comment on lines +558 to +564
Copy link
Collaborator

Choose a reason for hiding this comment

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

Part of me (albeit a part with limited convictions) feels like there HAS to be a better way to handle this. HF seems to be managing it internally with their pipelines, but this seems fine for the time being.

Copy link
Collaborator

Choose a reason for hiding this comment

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

idk man the format uses \n as a special char but doesn't seem to do mush escaping, many bets are off

Another option could be to do \n-based escaping and parsing of messages line by line, assuming [A-Z][a-z]+: is a metadata line

My experience also is that things like prefix_prompt get mangled in other code sometimes. I don't have great suggestions for how to check this well, other than:

  • continue doing qualitative reviews of prompt:output pairs for suspicious items;
  • migrate to a data format with clearer separation between data and metadata;
  • badger HF for a reference implementation;
  • give notice to users that HF Chat is "janky".


if not self.deprefix_prompt:
return text_output
else:
return [re.sub("^" + re.escape(prompt), "", i) for i in text_output]
return [re.sub("^" + re.escape(prefix_prompt), "", i) for i in text_output]


class LLaVA(Generator, HFCompatible):
Expand Down
30 changes: 30 additions & 0 deletions tests/generators/test_huggingface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,21 @@ def test_pipeline(hf_generator_config):
assert isinstance(item, str)


def test_pipeline_chat(mocker, hf_generator_config):
# uses a ~350M model with chat support
g = garak.generators.huggingface.Pipeline(
"microsoft/DialoGPT-small", config_root=hf_generator_config
)
mock_format = mocker.patch.object(
g, "_format_chat_prompt", wraps=g._format_chat_prompt
)
output = g.generate("Hello world!")
mock_format.assert_called_once()
assert len(output) == 1
for item in output:
assert isinstance(item, str)


def test_inference(mocker, hf_mock_response, hf_generator_config):
model_name = "gpt2"
mock_request = mocker.patch.object(
Expand Down Expand Up @@ -121,6 +136,21 @@ def test_model(hf_generator_config):
assert item is None # gpt2 is known raise exception returning `None`


def test_model_chat(mocker, hf_generator_config):
# uses a ~350M model with chat support
g = garak.generators.huggingface.Model(
"microsoft/DialoGPT-small", config_root=hf_generator_config
)
mock_format = mocker.patch.object(
g, "_format_chat_prompt", wraps=g._format_chat_prompt
)
output = g.generate("Hello world!")
mock_format.assert_called_once()
assert len(output) == 1
for item in output:
assert isinstance(item, str)


def test_select_hf_device():
from garak.generators.huggingface import HFCompatible
import torch
Expand Down
Loading