Adding modeling.py file

modeling.py

@@ -8,6 +8,7 @@ from .gpt import GPTBase
 from .aux_losses import entropy_reg, load_balancing_loss, router_z_loss
 from typing import Optional, List
 from dataclasses import dataclass
+import tiktoken
 
 
 @dataclass
@@ -37,8 +38,8 @@ class MoLM(PreTrainedModel):
         
         # Number of experts
         self.num_experts = config.num_experts
-        #print(f"Number of experts: {self.num_experts}")
-        #print(f"Expert configurations: {config.expert_configs}")
+        print(f"Number of experts: {self.num_experts}")
+        print(f"Expert configurations: {config.expert_configs}")
         assert len(config.expert_configs) == self.num_experts, "Number of expert configurations must match num_experts in config."
         self.expert_configs = config.expert_configs
 
@@ -52,6 +53,7 @@ class MoLM(PreTrainedModel):
 
         # Initialize experts using the provided configurations
         self.experts = nn.ModuleList([GPTBase(config=self.expert_configs[i]) for i in range(self.num_experts)])
+        self.tokenizer = tiktoken.get_encoding("gpt2")
         
         # Load pre-trained weights if provided
         if expert_weights is not None: