Skip to content

Commit

Permalink
Add missing arguments in pipeline parallel generate method
Browse files Browse the repository at this point in the history
Add two arguments: negative_prompt_ids and negative_prompt_attention_mask to generate method in pipeline_parallel.py.
These two arguments are available since transformers 4.32.0.
  • Loading branch information
notsyncing authored Sep 28, 2024
1 parent f71b38a commit 9604698
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions python/llm/src/ipex_llm/transformers/pipeline_parallel.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,8 @@ def generate(
synced_gpus: Optional[bool] = None,
assistant_model: Optional["PreTrainedModel"] = None,
streamer: Optional["BaseStreamer"] = None,
negative_prompt_ids: Optional[torch.Tensor] = None,
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
**kwargs,
):
if hasattr(self, 'pipeline_parallel_stages') and self.pipeline_parallel_stages > 1:
Expand Down Expand Up @@ -287,6 +289,8 @@ def generate(
synced_gpus=synced_gpus,
assistant_model=assistant_model,
streamer=streamer,
negative_prompt_ids=negative_prompt_ids,
negative_prompt_attention_mask=negative_prompt_attention_mask,
**kwargs)

GenerationMixin.generate = generate
Expand Down

0 comments on commit 9604698

Please sign in to comment.