Update 2 files

- /trainer.py
- /model.py

model.py

@@ -1,7 +1,6 @@
-#from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
-#from mamba_ssm.models.config_mamba import MambaConfig
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from mamba_ssm.models.config_mamba import MambaConfig
 
-from mamba import Mamba, ModelArgs
 
 
 import torch
@@ -13,8 +12,7 @@ class Model:
     def __init__(self, config: Config):
         self.__dict__ = dict(config.__dict__)
 
-        #self.model = MambaLMHeadModel(MambaConfig(**self.params.__dict__)).to(GetDevice())
-        self.model = Mamba(ModelArgs(**self.params.__dict__)).to(GetDevice())
+        self.model = MambaLMHeadModel(MambaConfig(**self.params.__dict__)).to(GetDevice())
         self.log()
 
 
@@ -45,57 +43,19 @@ class Model:
 
 
 
-    def generate_text(self, model,
-                tokenizer,
-                prompt: str,
-                n_tokens_to_gen: int = 50,
-                sample: bool = True,
-                top_k: int = 40):
-        
-        model = self.model
-        model.eval()
-        
-        input_ids = tokenizer.encode(prompt)
-        
-        for token_n in range(n_tokens_to_gen):
-            with torch.no_grad():
-                indices_to_input = input_ids
-                next_token_logits = model(indices_to_input)[:, -1]
-            
-            probs = F.softmax(next_token_logits, dim=-1)
-            (batch, vocab_size) = probs.shape
-            
-            if top_k is not None:
-                (values, indices) = torch.topk(probs, k=top_k)
-                probs[probs < values[:, -1, None]] = 0
-                probs = probs / probs.sum(axis=1, keepdims=True)
-            
-            if sample:
-                next_indices = torch.multinomial(probs, num_samples=1)
-            else:
-                next_indices = torch.argmax(probs, dim=-1)[:, None]
-            
-            input_ids = torch.cat([input_ids, next_indices], dim=1)
-
-        output_completions = [tokenizer.decode(output.tolist()) for output in input_ids][0]
-        
-        return output_completions
-
-
-    """
-    def generate_text(self, tokenizer, seed_text, num_predict):
+    def generate_text(self, seed_text, num_predict):
         max_len = num_predict + len(seed_text)
 
         with torch.no_grad():
-            encoded_ids = tokenizer.encode(seed_text)
+            encoded_ids = self.tokenizer.encode(seed_text)
             input_ids = torch.tensor(encoded_ids).unsqueeze(0).to(GetDevice())
             output = self.model.generate(input_ids, max_length=max_len)
 
             logits = output[0].tolist()
-            text = tokenizer.decode(logits)
+            text = self.tokenizer.decode(logits)
             
         return text
-    """
+
 
 
     @staticmethod

trainer.py

@@ -15,8 +15,8 @@ class Trainer:
         args = {'epoch': self.epoch, 'batch': self.batch, 'loss': loss}
         self.wandb(args)
 
-        #if self.batch % 200 == 0:
-        #    print(f'{self.model.generate_text(self.model.tokenizer, self.inference.seed_text, self.inference.n_predict)}')
+        if self.batch % 200 == 0:
+            print(f'{self.model.generate_text(self.inference.seed_text, self.inference.n_predict)}')
 
 
     def train(self, batches):