Add TFAutoModelForSemanticSegmentation to the main __init__.py (h…

…uggingface#18600) Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
oneraghavan · Sep 26, 2022 · 4c5b6d2 · 4c5b6d2
1 parent 8b85df1
commit 4c5b6d2
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 0 deletions.
diff --git a/docs/source/en/model_doc/auto.mdx b/docs/source/en/model_doc/auto.mdx
@@ -186,6 +186,10 @@ Likewise, if your `NewModel` is a subclass of [`PreTrainedModel`], make sure its
 
 [[autodoc]] TFAutoModelForImageClassification
 
+## TFAutoModelForSemanticSegmentation
+
+[[autodoc]] TFAutoModelForSemanticSegmentation
+
 ## TFAutoModelForMaskedLM
 
 [[autodoc]] TFAutoModelForMaskedLM

diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py
@@ -2105,6 +2105,7 @@
             "TFAutoModelForNextSentencePrediction",
             "TFAutoModelForPreTraining",
             "TFAutoModelForQuestionAnswering",
+            "TFAutoModelForSemanticSegmentation",
             "TFAutoModelForSeq2SeqLM",
             "TFAutoModelForSequenceClassification",
             "TFAutoModelForSpeechSeq2Seq",
@@ -4599,6 +4600,7 @@
             TFAutoModelForNextSentencePrediction,
             TFAutoModelForPreTraining,
             TFAutoModelForQuestionAnswering,
+            TFAutoModelForSemanticSegmentation,
             TFAutoModelForSeq2SeqLM,
             TFAutoModelForSequenceClassification,
             TFAutoModelForSpeechSeq2Seq,

diff --git a/src/transformers/models/auto/__init__.py b/src/transformers/models/auto/__init__.py
@@ -128,6 +128,7 @@
         "TFAutoModelForNextSentencePrediction",
         "TFAutoModelForPreTraining",
         "TFAutoModelForQuestionAnswering",
+        "TFAutoModelForSemanticSegmentation",
         "TFAutoModelForSeq2SeqLM",
         "TFAutoModelForSequenceClassification",
         "TFAutoModelForSpeechSeq2Seq",
@@ -271,6 +272,7 @@
             TFAutoModelForNextSentencePrediction,
             TFAutoModelForPreTraining,
             TFAutoModelForQuestionAnswering,
+            TFAutoModelForSemanticSegmentation,
             TFAutoModelForSeq2SeqLM,
             TFAutoModelForSequenceClassification,
             TFAutoModelForSpeechSeq2Seq,

diff --git a/src/transformers/utils/dummy_tf_objects.py b/src/transformers/utils/dummy_tf_objects.py
@@ -362,6 +362,13 @@ def __init__(self, *args, **kwargs):
         requires_backends(self, ["tf"])
 
 
+class TFAutoModelForSemanticSegmentation(metaclass=DummyObject):
+    _backends = ["tf"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["tf"])
+
+
 class TFAutoModelForSeq2SeqLM(metaclass=DummyObject):
     _backends = ["tf"]