Update README.md

README.md

@@ -18,25 +18,50 @@ The official code repository is available at: https://github.com/selmiss/MuMo
 
 - **Model Type**: MuMo Pretrained Model
 - **Training Data**: Molecular structures and properties
-- **Framework**: PyTorch + Transformers
+- **Framework**: PyTorch + Transformers + Mamba-ssm
 
 ## Usage
 
+### Loading the Model
+
+MuMo uses a custom loading function. Here's how to load the pretrained model:
+
+```shell
+git clone https://github.com/selmiss/MuMo.git
+```
+
 ```python
-from transformers import AutoConfig, AutoTokenizer, AutoModel
-
-# Load model
-model_path = "zihaojing/mumo-pretrain"
-config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
-
-# Example usage
-smiles = "CCO"  # Ethanol
-inputs = tokenizer(smiles, return_tensors="pt")
-outputs = model(**inputs)
+from transformers import AutoConfig, AutoTokenizer
+from model.load_model import load_model
+from dataclasses import dataclass
+
+# Load configuration and tokenizer
+repo = "zihaojing/MuMo-Pretrained"
+config = AutoConfig.from_pretrained(repo, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(repo)
+
+# Set up model arguments
+@dataclass
+class ModelArgs:
+    model_name_or_path: str = repo
+    model_class: str = "MuMoFinetune"  # or "MuMoPretrain" for pretraining
+    cache_dir: str = None
+    model_revision: str = "main"
+    use_auth_token: bool = False
+    task_type: str = None  # e.g., "classification" or "regression" for finetuning
+
+model_args = ModelArgs()
+
+# Load the model
+model = load_model(config, tokenizer=tokenizer, model_args=model_args)
 ```
 
+**Notes:**
+- Use `model_class="MuMoPretrain"` for pretraining or inference
+- Use `model_class="MuMoFinetune"` for finetuning tasks
+- Set `task_type` to `"classification"` or `"regression"` when using `MuMoFinetune`
+- The model supports loading from both Hugging Face Hub (e.g., `"zihaojing/MuMo-Pretrained"`) and local paths (e.g., `"/path/to/model"`)
+
 ## Training Details
 
 - Training script: See the [official GitHub repository](https://github.com/selmiss/MuMo) for details.