Update 2 files

- /model.py
- /trainer.py

model.py

@@ -9,16 +9,28 @@ class Model:
         self.__dict__ = dict(config.__dict__)
 
         self.model = MambaLMHeadModel(MambaConfig(**self.params.__dict__)).to(GetDevice())
-        self.Log()
+        self.log()
 
 
-    def Log(self):
+    def log(self):
         model_size, rounded_model_size = GetNumParams(self.model)
         print(f"Model has {model_size} ({rounded_model_size}) parameters")
         print(f"Model's embedding size is {self.params.vocab_size}")
 
 
-    def AutoRegressiveLossFunction(self, input_ids, labels=None, criterion=None):
+    def parameters():
+        return self.model.parameters()
+
+
+    def unfreeze():
+        self.model.train()
+
+
+    def freeze():
+        self.model.eval()
+
+
+    def compute_loss(self, input_ids, labels=None, criterion=None):
         lm_logits = self.model(input_ids).logits
 
         labels = input_ids.to(self.model.device)
@@ -31,7 +43,7 @@ class Model:
         return lm_loss
 
 
-    def GenerateText(self, tokenizer, seed_text, num_predict):
+    def generate_text(self, tokenizer, seed_text, num_predict):
         max_len = num_predict + len(seed_text)
 
         with torch.no_grad():
@@ -45,5 +57,5 @@ class Model:
 
 
     @staticmethod
-    def SavePretrained(self, path='./'):
+    def save_pretrained(self, path='./'):
         self.model.save_pretrained(path)

trainer.py

@@ -6,13 +6,34 @@ from model import Model
 
 class Trainer:
     def __init__(self, config: Config):
-        self.__dict__ = dict(config.__dict__)
-
-        #self.wandb = Wandb(config.wandb)
+        self.__dict__ = dict(config.trainer.__dict__)
 
         
+        #self.wandb = Wandb(config.wandb)
         self.model = Model(config.model)
         
+        self.optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)
+
+
+    def log(self, loss: float):
+        print(f"Epoch: {self.epoch} / {self.num_epochs}\t\tBatch: {self.batch} / {self.num_batches}\t\tLoss: {round(loss, 4)}")
+        
+
+
+    def train(self, batches): 
+        #pass
+
+        model.unfreeze()
+
+        for self.epoch in range(self.num_epochs):
+            for self.batch in range(self.num_batches):
+                ids = batches[batch]
+
+                loss = model.compute_loss(ids)
+
+                self.optimizer.zero_grad()
+                loss.backward()
+                self.optimizer.step()
 
-    def train(self, dataset): # TODO: Implement
-        pass
+                self.log(loss.item())
+                #Train.LogStep(infer_config, log_config, epoch, num_epochs, batch, num_batches, loss)