Update 3 files

- /util.py
- /trainer.py
- /model.py

model.py

@@ -0,0 +1,24 @@
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from mamba_ssm.models.config_mamba import MambaConfig
+
+from util import Config, GetDevice
+
+
+class Model:
+    def __init__(self, config: Config):
+        self.__dict__ = dict(config.__dict__)
+
+        self.model = MambaLMHeadModel(MambaConfig(params)).to(GetDevice())
+
+
+    def AutoRegressiveLossFunction(self, input_ids, labels=None, criterion=None):
+        lm_logits = self.model(input_ids).logits
+
+        labels = input_ids.to("cuda")
+        shift_logits = lm_logits[:, :-1, :].contiguous()
+        labels = labels[:, 1:].contiguous()
+
+        loss_fct = criterion or torch.nn.CrossEntropyLoss()
+        lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
+
+        return lm_loss

trainer.py

@@ -1,5 +1,7 @@
 from util import Config
+
 from logger import Wandb
+from model import Model
 
 
 class Trainer:
@@ -7,4 +9,5 @@ class Trainer:
         self.__dict__ = dict(config.__dict__)
 
         self.wandb = Wandb(config.wandb)
+        self.model = Model(config.model)
         

util.py

@@ -3,6 +3,11 @@ import math
 import random
 
 
+
+def GetDevice():
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
 def RandomCode():
     code = '';
     chars = '0123456789abcdef'