Adding modeling.py file

modeling.py

@@ -107,7 +107,7 @@ class MoLM(PreTrainedModel):
                 # Expand the expert_mask to match the logits shape (batch_size, 1, 1)
                 expert_mask_expanded = expert_mask.unsqueeze(-1).unsqueeze(-1).float()
 
-                expert_output = expert(input_ids, targets=targets, date=date, get_logits=True, **kwargs)
+                expert_output = expert(input_ids, targets=targets, date=date, **kwargs)
 
                 logits = expert_output["logits"]
                 loss_to_log = expert_output["loss_to_log"]
@@ -234,10 +234,10 @@ class MoLM(PreTrainedModel):
     def generate_from_string(self, in_str, max_new_tokens, date=None, temperature=1.0, top_k=None):
         idx = (
             torch.tensor(
-                self.tokenizer.encode(in_str, allowed_special={"<|endoftext|>"})
+                self.tokenizer.encode(in_str)
             )
             .view(1, -1)
-            .to(self.lm_head.weight.device)
+            .to(next(self.parameters()).device)
         )
         out_idx = (
             self.generate(idx, max_new_tokens, date, temperature, top_k)