Adding OneFormer model to DeepChem (deepchem#4146)

* Added OneFormer model + tests + init + docs

* Modified docs

* Added model to rst

* Removed dummy tester file

* Added save/reload + overfit tests

* Try 1 DQC Fix

* Try 2 DQC Fix

* Made changes to tests

* Modified mini config

* Updated tests

deepchem/models/__init__.py

@@ -57,6 +57,7 @@
     from deepchem.models.torch_models import HuggingFaceModel
     from deepchem.models.torch_models import Chemberta
     from deepchem.models.torch_models import MoLFormer
+    from deepchem.models.torch_models import OneFormer
 except ImportError as e:
     logger.warning(e)
 

deepchem/models/torch_models/__init__.py

@@ -51,6 +51,7 @@
     from deepchem.models.torch_models.chemberta import Chemberta
     from deepchem.models.torch_models.molformer import MoLFormer
     from deepchem.models.torch_models.prot_bert import ProtBERT
+    from deepchem.models.torch_models.oneformer import OneFormer
 
 except ModuleNotFoundError as e:
     logger.warning(f'Skipped loading modules with transformers dependency. {e}')

deepchem/models/torch_models/hf_models.py

@@ -11,7 +11,7 @@
 from deepchem.trans import Transformer, undo_transforms
 from deepchem.utils.typing import LossFn, OneOrMany
 from transformers.data.data_collator import DataCollatorForLanguageModeling
-from transformers.models.auto import AutoModel, AutoModelForSequenceClassification, AutoModelForMaskedLM
+from transformers.models.auto import AutoModel, AutoModelForSequenceClassification, AutoModelForMaskedLM, AutoModelForUniversalSegmentation
 
 logger = logging.getLogger(__name__)
 
@@ -216,6 +216,9 @@ def load_from_pretrained(  # type: ignore
             elif self.task in ['mtr', 'regression', 'classification']:
                 self.model = AutoModelForSequenceClassification.from_pretrained(
                     model_dir)
+            elif self.task == "universal_segmentation":
+                self.model = AutoModelForUniversalSegmentation.from_pretrained(
+                    model_dir)
             else:
                 self.model = AutoModel.from_pretrained(model_dir)
         elif not from_hf_checkpoint: