riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-03 03:31:05 +06:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

transformers.fx.symbolic_trace supports inputs_embeds (#31574)

Browse Source 
* symbolic trace supports inputs_embeds

* fix test?

* Update tests/test_modeling_common.py

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

---------

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>
This commit is contained in:
fxmarty 2024-07-08 13:17:28 +02:00 committed by GitHub
parent e5ca9b057c
commit ba743700f4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/transformers/utils/fx.py
View File

@ -995,6 +995,13 @@ class HFTracer(Tracer):
inputs_dict[input_name] = torch.zeros(
*shape, model.config.input_feat_per_channel, dtype=torch.float, device=device
)
elif "inputs_embeds" in input_name:
batch_size = shape[0]
sequence_length = shape[-1]
inputs_dict[input_name] = torch.zeros(
batch_size, sequence_length, model.config.hidden_size, dtype=torch.float, device=device
)
elif "visual_feats" in input_name:
inputs_dict[input_name] = torch.zeros(
shape

16
tests/test_modeling_common.py
View File

@ -1158,6 +1158,7 @@ class ModelTesterMixin:
"input_features",
"input_ids",
"input_values",
"inputs_embeds",
"pixel_values",
"token_type_ids",
"visual_feats",
@ -1214,16 +1215,27 @@ class ModelTesterMixin:
(past_mask, inputs_to_test[1]["attention_mask"]), dim=1
)
if "inputs_embeds" in inspect.signature(model.forward).parameters:
inputs_to_test.append(
{
"inputs_embeds": torch.rand(
2, 2, model.config.hidden_size, dtype=torch.float, device=torch_device
)
}
)
for inps in inputs_to_test:
filtered_inputs = {k: v for (k, v) in inps.items() if k in input_names}
input_names = list(filtered_inputs.keys())
input_names_to_trace = list(filtered_inputs.keys())
if model.__class__.__name__ in set(MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES.values()) and (
not hasattr(model.config, "problem_type") or model.config.problem_type is None
):
model.config.problem_type = "single_label_classification"
traced_model = symbolic_trace(model, input_names)
model.config.use_cache = "past_key_values" in input_names_to_trace
traced_model = symbolic_trace(model, input_names_to_trace)
with torch.no_grad():
traced_output = traced_model(**filtered_inputs)