fix imports

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

Hide file tree

Showing 2 changed files with 1 addition and 8 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/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)