make multimolecule.module for shared modules

Signed-off-by: Zhiyuan Chen <[email protected]>

multimolecule/__init__.py

@@ -8,6 +8,10 @@
     UtrLmForCrisprOffTarget,
 )
 from .models import (
+    BaseHeadConfig,
+    HeadConfig,
+    MaskedLMHeadConfig,
+    PretrainedConfig,
     RnaBertConfig,
     RnaBertForMaskedLM,
     RnaBertForNucleotideClassification,
@@ -54,6 +58,10 @@
 from .tokenisers import RnaTokenizer
 
 __all__ = [
+    "PretrainedConfig",
+    "BaseHeadConfig",
+    "HeadConfig",
+    "MaskedLMHeadConfig",
     "models",
     "tokenisers",
     "RnaTokenizer",

multimolecule/downstream/crispr_off_target.py

@@ -7,13 +7,13 @@
 from torch import Tensor
 from transformers.modeling_outputs import ModelOutput
 
-from multimolecule.models.modeling_utils import ClassificationHead, Criterion
 from multimolecule.models.rnabert import RnaBertConfig, RnaBertModel, RnaBertPreTrainedModel
 from multimolecule.models.rnafm import RnaFmConfig, RnaFmModel, RnaFmPreTrainedModel
 from multimolecule.models.rnamsm import RnaMsmConfig, RnaMsmModel, RnaMsmPreTrainedModel
 from multimolecule.models.splicebert import SpliceBertConfig, SpliceBertModel, SpliceBertPreTrainedModel
 from multimolecule.models.utrbert import UtrBertConfig, UtrBertModel, UtrBertPreTrainedModel
 from multimolecule.models.utrlm import UtrLmConfig, UtrLmModel, UtrLmPreTrainedModel
+from multimolecule.module import ClassificationHead, Criterion
 
 
 class RnaBertForCrisprOffTarget(RnaBertPreTrainedModel):

multimolecule/models/__init__.py

@@ -1,4 +1,5 @@
 from ..tokenisers.rna import RnaTokenizer
+from .configuration_utils import BaseHeadConfig, HeadConfig, MaskedLMHeadConfig, PretrainedConfig
 from .rnabert import (
     RnaBertConfig,
     RnaBertForMaskedLM,
@@ -55,6 +56,10 @@
 )
 
 __all__ = [
+    "PretrainedConfig",
+    "HeadConfig",
+    "BaseHeadConfig",
+    "MaskedLMHeadConfig",
     "RnaTokenizer",
     "RnaBertConfig",
     "RnaBertModel",

multimolecule/models/configuration_utils.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
+from collections import OrderedDict
 from dataclasses import asdict, dataclass, is_dataclass
 from typing import Optional
 
 from transformers.configuration_utils import PretrainedConfig as _PretrainedConfig
 
 
 class PretrainedConfig(_PretrainedConfig):
-    head: HeadConfig
+    head: BaseHeadConfig
 
     def __init__(
         self, pad_token_id=0, bos_token_id=1, eos_token_id=2, unk_token_id=3, mask_token_id=4, null_token_id=5, **kwargs
@@ -38,8 +39,12 @@ def to_dict(self):
         return output
 
 
+class HeadConfig(OrderedDict):
+    pass
+
+
 @dataclass
-class HeadConfig:
+class BaseHeadConfig(HeadConfig):
     r"""
     This is the configuration class to store the configuration of a prediction head. It is used to instantiate a
     prediction head according to the specified arguments, defining the head architecture.
@@ -82,7 +87,7 @@ class HeadConfig:
 
 
 @dataclass
-class MaskedLMHeadConfig:
+class MaskedLMHeadConfig(HeadConfig):
     r"""
     This is the configuration class to store the configuration of a prediction head. It is used to instantiate a
     prediction head according to the specified arguments, defining the head architecture.