Upload model

chatNT.py

@@ -588,6 +588,19 @@ class TorchMultiOmicsModel(PreTrainedModel):
     config_class = ChatNTConfig
 
     def __init__(self, config: ChatNTConfig) -> None:
+        if isinstance(config, dict):
+            # If config is a dictionary instead of ChatNTConfig (which can happen
+            # depending how the config was saved), we convert it to the config
+            config["gpt_config"]["rope_config"] = RotaryEmbeddingConfig(
+                **config["gpt_config"]["rope_config"]
+            )
+            config["gpt_config"] = GptConfig(**config["gpt_config"])
+            config["esm_config"] = ESMTransformerConfig(**config["esm_config"])
+            config["perceiver_resampler_config"] = PerceiverResamplerConfig(
+                **config["perceiver_resampler_config"]
+            )
+            config = ChatNTConfig(**config)  # type: ignore
+
         super().__init__(config=config)
         self.gpt_config = config.gpt_config
         self.esm_config = config.esm_config