Uploading model.pt

model.py

@@ -1,6 +1,13 @@
 import torch
 from torch import nn
-from transformers import PreTrainedModel, AutoConfig, AutoModel
+from transformers import PreTrainedModel, PretrainedConfig , AutoModel
+
+class EmbeddingMoEConfig(PretrainedConfig):
+    def __init__(self, output_dim=128, num_experts=2, dropout_rate=0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.output_dim = output_dim
+        self.num_experts = num_experts
+        self.dropout_rate = dropout_rate
 
 
 # Expert class using pre-trained BERT
@@ -69,7 +76,7 @@ class GatingNetwork(nn.Module):
 
 # Mixture of Experts for sentence embeddings using BERT
 class EmbeddingMoE(PreTrainedModel):
-    config_class = AutoConfig
+    config_class = EmbeddingMoEConfig
 
     def __init__(self, config):
         super().__init__(config)