Update README.md

README.md

@@ -19,7 +19,60 @@ We are releasing the weights for the three best-performing models:
 
 ## Usage
 
+Below are sample snippets for loading each model with its pretrained weights. Once loaded, you can run inference by passing your pre-processed ECG data to the model’s forward method. Please refer to the project’s GitHub repository for end-to-end examples demonstrating how to use these models.
 
+### S4-ECG
 
-## Citation
+```bash
+import torch
+import lightning.pytorch as pl
+from src.lightning import S4Model
+
+def load_from_checkpoint(pl_model, checkpoint_path):
+    """ load from checkpoint function that is compatible with S4
+    """
+    lightning_state_dict = torch.load(checkpoint_path)
+    state_dict = lightning_state_dict["state_dict"]
+    
+    for name, param in pl_model.named_parameters():
+        param.data = state_dict[name].data
+    for name, param in pl_model.named_buffers():
+        param.data = state_dict[name].data
+
+checkpoint_path = "path/to/your/benchmark_acs_state_v0/dmaxlwcg/checkpoints/epoch=49-step=1100.ckpt"
+
+model = S4Model(init_lr=1e-4,
+                d_input=3,
+                d_output=1)
+
+load_from_checkpoint(model, checkpoint_path)
+```
+
+### ResNet-18
 
+```bash
+import torch
+import lightning.pytorch as pl
+from src.lightning import ResNet18_1D
+
+checkpoint_path = "path/to/your/benchmark_acs_resnet18_1d_final_vf/2vud5fft/checkpoints/epoch=37-step=418.ckpt"
+model = ResNet18_1D.load_from_checkpoint(checkpoint_path)
+```
+
+### HuBERT-ECG
+
+```bash
+import torch
+from hubert_ecg import HuBERTECG, HuBERTECGConfig
+from hubert_ecg_classification import HuBERTForECGClassification
+
+path = "path/to/your/hubert_3_iteration_300_finetuned_simdmsnv.pt"
+checkpoint = torch.load(path, map_location='cpu')
+config = checkpoint['model_config']
+hubert_ecg = HuBERTECG(config)
+hubert_ecg = HuBERTForECGClassification(hubert_ecg)
+hubert_ecg.load_state_dict(checkpoint['model_state_dict'])
+```
+
+## Citation
+*Paper currently under review*