Update file model.py

model.py

@@ -44,6 +44,45 @@ class Model:
         return lm_loss
 
 
+
+    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(prompt, return_tensors='pt').input_ids
+        
+        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):
         max_len = num_predict + len(seed_text)
 
@@ -56,6 +95,7 @@ class Model:
             text = tokenizer.decode(logits)
             
         return text
+    """
 
 
     @staticmethod