Update tiny models for pipeline testing. (#24364)

* fix

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2023-06-20 14:43:10 +02:00
committed by GitHub
parent 56efbf4301
commit c23d131eab
15 changed files with 110 additions and 29 deletions

View File

@@ -383,11 +383,22 @@ class GitModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin,
all_model_classes = (GitModel, GitForCausalLM) if is_torch_available() else ()
all_generative_model_classes = (GitForCausalLM,) if is_torch_available() else ()
pipeline_model_mapping = (
{"feature-extraction": GitModel, "text-generation": GitForCausalLM} if is_torch_available() else {}
{"feature-extraction": GitModel, "image-to-text": GitForCausalLM, "text-generation": GitForCausalLM}
if is_torch_available()
else {}
)
fx_compatible = False
test_torchscript = False
# `GitForCausalLM` doesn't fit into image-to-text pipeline. We might need to overwrite its `generate` function.
def is_pipeline_test_to_skip(
self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name
):
if pipeline_test_casse_name == "ImageToTextPipelineTests":
return True
return False
# special case for GitForCausalLM model
def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)