Update chatNT.py

chatNT.py

@@ -364,21 +364,21 @@ class ChatNTConfig(PretrainedConfig):
         return output
 
 
-class TorchBioBrainDecoder(nn.Module):
+class ChatNTDecoder(nn.Module):
     def __init__(
         self,
         gpt_config: GptConfig,
         seq_token_id: int,
     ):
         """
-        Initializes the BioBrain decoder, using a GPT model for text generation with
+        Initializes the ChatNT decoder, using a GPT model for text generation with
         bio embeddings.
 
         Args:
             gpt_config: Configuration for the GPT model
             seq_token_id: Index of the SEQ token
         """
-        super(TorchBioBrainDecoder, self).__init__()
+        super(ChatNTDecoder, self).__init__()
         self.gpt_config = gpt_config
         self.seq_token_id = seq_token_id
 
@@ -582,7 +582,7 @@ class TorchBioBrainDecoder(nn.Module):
         return logits_acc, tokens_acc
 
 
-class TorchMultiOmicsModel(PreTrainedModel):
+class ChatNT(PreTrainedModel):
     config_class = ChatNTConfig
 
     def __init__(self, config: ChatNTConfig) -> None:
@@ -625,11 +625,11 @@ class TorchMultiOmicsModel(PreTrainedModel):
         # Correct seq_token_id
         self.seq_token_id -= 1
 
-        self.biobrain_encoder = TorchBioBrainEncoder(nt_config=self.nt_config)
-        self.biobrain_decoder = TorchBioBrainDecoder(
+        self.chatnt_encoder = ChatNTEncoder(nt_config=self.nt_config)
+        self.chatnt_decoder = ChatNTDecoder(
             gpt_config=self.gpt_config, seq_token_id=self.seq_token_id
         )
-        self.projection_model = TorchMultiModalPerceiverResamplerProjection(
+        self.projection_model = MultiModalPerceiverResamplerProjection(
             perceiver_resampler_config=self.perceiver_resampler_config,
             input_embed_dim=self.nt_config.embed_dim,
             embed_dim=self.gpt_config.embed_dim,
@@ -702,7 +702,7 @@ class TorchMultiOmicsModel(PreTrainedModel):
             if projected_bio_embeddings is None:
                 # Compute bio sequences embeddings
                 bio_embeddings_list = [
-                    self.biobrain_encoder(bio_token_ids=bio_token_ids[:, bio_seq_num])
+                    self.chatnt_encoder(bio_token_ids=bio_token_ids[:, bio_seq_num])
                     for bio_seq_num in range(num_bio_sequences)
                 ]
 
@@ -718,7 +718,7 @@ class TorchMultiOmicsModel(PreTrainedModel):
                 projected_bio_embeddings = torch.stack(projected_bio_embeddings, dim=1)
 
         # decode
-        logits = self.biobrain_decoder(
+        logits = self.chatnt_decoder(
             english_token_ids=english_token_ids,
             projected_bio_embeddings=projected_bio_embeddings,
         )
@@ -1498,12 +1498,12 @@ class RobertaLMHead(nn.Module):
         return {"embeddings": embeddings, "logits": logits}
 
 
-class TorchNucleotideTransformer(nn.Module):
+class NucleotideTransformer(nn.Module):
     def __init__(
         self,
         nt_config: NucleotideTransformerConfig,
     ):
-        super(TorchNucleotideTransformer, self).__init__()
+        super(NucleotideTransformer, self).__init__()
         self.nt_config = nt_config
 
         # Other cases are not implemented
@@ -1599,14 +1599,14 @@ def build_padding_attention_mask(
     return padding_mask
 
 
-class TorchBioBrainEncoder(nn.Module):
+class ChatNTEncoder(nn.Module):
     def __init__(
         self,
         nt_config: NucleotideTransformerConfig,
     ):
-        super(TorchBioBrainEncoder, self).__init__()
+        super(ChatNTEncoder, self).__init__()
         self.nt_config = nt_config
-        self.nt_model = TorchNucleotideTransformer(self.nt_config)
+        self.nt_model = NucleotideTransformer(self.nt_config)
 
     def forward(
         self,
@@ -1626,7 +1626,7 @@ class TorchBioBrainEncoder(nn.Module):
         return bio_embeddings
 
 
-class TorchMultiModalPerceiverResamplerBlock(nn.Module):
+class MultiModalPerceiverResamplerBlock(nn.Module):
     def __init__(
         self,
         num_heads: int,
@@ -1714,7 +1714,7 @@ class TorchMultiModalPerceiverResamplerBlock(nn.Module):
         return {"embeddings": x}
 
 
-class TorchMultiModalPerceiverResampler(nn.Module):
+class MultiModalPerceiverResampler(nn.Module):
     """
     Perceiver Resampler model, made of successive PerceiverResamplerBlocks.
     """
@@ -1736,7 +1736,7 @@ class TorchMultiModalPerceiverResampler(nn.Module):
         self.name = name
         self.layers = nn.ModuleList(
             [
-                TorchMultiModalPerceiverResamplerBlock(
+                MultiModalPerceiverResamplerBlock(
                     num_heads=self.config.attention_heads,
                     embed_dim=self.config.embed_dim,
                     key_size=self.config.key_size,
@@ -1823,7 +1823,7 @@ class TorchMultiModalPerceiverResampler(nn.Module):
         return outs
 
 
-class TorchMultiModalPerceiverResamplerProjection(nn.Module):
+class MultiModalPerceiverResamplerProjection(nn.Module):
     def __init__(
         self,
         perceiver_resampler_config: PerceiverResamplerConfig,
@@ -1843,7 +1843,7 @@ class TorchMultiModalPerceiverResamplerProjection(nn.Module):
 
         self.bio_projection = nn.Linear(input_embed_dim, embed_dim)
         self.token_embedding = nn.Embedding(english_vocab_size, embed_dim)
-        self.perceiver_resampler = TorchMultiModalPerceiverResampler(config=self.config)
+        self.perceiver_resampler = MultiModalPerceiverResampler(config=self.config)
 
     def forward(
         self,