fix imports

Signed-off-by: Zhiyuan Chen <[email protected]>
DLS5-Omics · Mar 26, 2024 · 0f62c93 · 0f62c93
1 parent 3188088
commit 0f62c93
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 10 deletions.
diff --git a/multimolecule/__init__.py b/multimolecule/__init__.py
@@ -1,10 +1,3 @@
-from transformers import AutoTokenizer
-
 from . import models, tokenizers
-from .models import RnaBertConfig
-from .tokenizers import RnaTokenizer
-
-AutoTokenizer.register(RnaBertConfig, RnaTokenizer)
-
 
 __all__ = ["models", "tokenizers"]
diff --git a/multimolecule/models/__init__.py b/multimolecule/models/__init__.py
@@ -1,3 +1,3 @@
-from .rnabert import RnaBertConfig, RnaBertModel
+from .rnabert import RnaBertConfig, RnaBertModel, RnaTokenizer
 
-__all__ = ["RnaBertConfig", "RnaBertModel"]
+__all__ = ["RnaBertConfig", "RnaBertModel", "RnaTokenizer"]
diff --git a/multimolecule/models/rnabert/__init__.py b/multimolecule/models/rnabert/__init__.py
@@ -5,7 +5,7 @@
 from .configuration_rnabert import RnaBertConfig
 from .modeling_rnabert import RnaBertModel
 
-__all__ = ["RnaBertConfig", "RnaBertModel"]
+__all__ = ["RnaBertConfig", "RnaBertModel", "RnaTokenizer"]
 
 AutoConfig.register("rnabert", RnaBertConfig)
 AutoModel.register(RnaBertConfig, RnaBertModel)