Update modeling_adaptor.py

modeling_adaptor.py

@@ -1,7 +1,11 @@
+import os
+import json
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+
 class MoAGate(nn.Module):
     def __init__(self, num_adaptors, hidden_dim):
         super().__init__()
@@ -84,4 +88,16 @@ class MixtureOfAdaptors(nn.Module):
             )
         return adaptor_cache
 
+    @classmethod
+    def load(cls, input_path):
+        with open(os.path.join(input_path, "config.json")) as fIn:
+            config = json.load(fIn)
+
+        adaptor = cls(**config)
+        adaptor.load_state_dict(
+            torch.load(
+                os.path.join(input_path, "adaptor.pth"), weights_only=True
+            )
+        )
+        return adaptor