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

Jump to bottom

[Text Generation][MPT] Filter out the appropriate engine_inputs for autoregressive_inference #1151

Merged
merged 3 commits into from
Aug 1, 2023
Merged

[Text Generation][MPT] Filter out the appropriate engine_inputs for autoregressive_inference #1151

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
94f72c7
initial commit
dbogunowicz Jul 28, 2023
55c3e5b
Merge branch 'main' into fix/damian/positions
dbogunowicz Aug 1, 2023
a49b411
Merge branch 'main' into fix/damian/positions
dbogunowicz Aug 1, 2023
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
13 changes: 12 additions & 1 deletion src/deepsparse/transformers/pipelines/text_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,7 @@ def autoregressive_inference(
:return: The new, generated token and the logits for the new token
(with dimensions ['batch_size', 'num_tokens', 'vocab_size'])
"""

new_token = tokens[-1]
# padding is added to left, so attention mask is 1s from the
# right up to the number of total tokens (prompt + generated)
Expand All @@ -444,7 +445,17 @@ def autoregressive_inference(
positions -= 1
input_ids = numpy.array([[new_token]])
causal_mask = create_causal_mask(input_ids, attention_mask)
engine_inputs = [input_ids, attention_mask, positions, causal_mask]

# filter out the inputs that are not needed by the engine
engine_inputs_map = dict(
input_ids=input_ids,
attention_mask=attention_mask,
causal_mask=causal_mask,
positions=positions,
)
engine_inputs = [
engine_inputs_map[name] for name in self.engine.onnx_input_names_no_cache
]

generated_token, generated_logits = self.engine(engine_inputs)

Expand Down
Loading