Update llm_trainer.py

llm_trainer.py

@@ -1,80 +1,55 @@
-import torch
-from torch.nn import functional as F
-from transformers import PreTrainedTokenizer, AutoTokenizer
-
-from llm_trainer.dataset.DataLoader import DataLoader
-
-class LLMTrainer:
-    def __init__(self,
-                 model: torch.nn.Module = None,
-                 tokenizer: PreTrainedTokenizer | AutoTokenizer = None,
-                 model_returns_logits: bool = False):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"Training on: {self.device}")
-
-        if tokenizer is None:
-            tokenizer = AutoTokenizer.from_pretrained("gpt2")
-
-        self.tokenizer = tokenizer
-
-        if model is None:
-            raise ValueError("Specify a model.")
-        self.model = model
-
-        self.train_loader = None
-        self.current_step: int = 0
-
-        self.model_returns_logits = model_returns_logits
-
-    def generate_text(self, prompt: str = "Once upon a time", n_return_sequences: int = 4, length: int = 32) -> None:
-        # Make sure the model is on the same device
-        self.model.to(self.device)
-        self.model.eval()
-
-        tokens = self.tokenizer.encode(prompt, return_tensors="pt").type(torch.long)
-        tokens = tokens.repeat(n_return_sequences, 1)
-
-        generated_tokens = tokens.to(self.device)
-        with torch.no_grad():
-            while generated_tokens.size(1) < length:
-
-                with torch.autocast(device_type=self.device, dtype=torch.bfloat16):
-                    if self.model_returns_logits:
-                        logits = self.model(generated_tokens)
-                    else:
-                        logits = self.model(generated_tokens).logits
-
-                # logits.shape = (batch_size, context_window, vocab_size)
-
-                logits = logits[:, -1, :]  # Get last token logits (B, vocab_size)
-                probs = F.softmax(logits, dim=-1)  # Convert to probabilities
-
-                # Top-k sampling
-                topk_probs, topk_indices = torch.topk(probs, k=10, dim=-1)
-                sampled_indices = torch.multinomial(topk_probs, 1)  # Shape: (B, 1)
-                next_tokens = torch.gather(topk_indices, -1, sampled_indices)  # (B, 1)
-
-                # Append generated token to sequence
-                generated_tokens = torch.cat((generated_tokens, next_tokens), dim=1)
-
-        # print the generated text
-        continuations = []
-        for i in range(n_return_sequences):
-            tokens = generated_tokens[i, :length].tolist()
-            decoded = self.tokenizer.decode(tokens)
-            print(f"=== sample {i} ===\n{decoded}")
-            continuations.append(decoded)
-        return continuations
-
-    def load_checkpoint(self, checkpoint_path: str) -> None:
-        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
-
-        # If the model was saved after running `torch.compile` then the names of its layers were changed.
-        # Need to change it back.
-        new_state_dict = {k.replace("_orig_mod.", ""): v for k, v in checkpoint['model_state_dict'].items()}
-        self.model.to(self.device)
-        self.model.load_state_dict(new_state_dict)
-        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-        self.train_loader: DataLoader = checkpoint["train_loader"]
-
-        self.current_step = checkpoint['step']  # Resume from the last step
+import torch
+from torch.nn import functional as F
+from transformers import PreTrainedTokenizer, AutoTokenizer
+
+class LLMTrainer:
+    def __init__(self,
+                 model: torch.nn.Module = None,
+                 tokenizer: PreTrainedTokenizer | AutoTokenizer = None,
+                 model_returns_logits: bool = False):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+        self.tokenizer = tokenizer
+        self.model = model
+
+        self.train_loader = None
+        self.current_step: int = 0
+
+        self.model_returns_logits = model_returns_logits
+
+    def generate_text(self, prompt: str = "Once upon a time", n_return_sequences: int = 4, length: int = 32) -> None:
+        self.model.to(self.device)
+        self.model.eval()
+
+        tokens = self.tokenizer.encode(prompt, return_tensors="pt").type(torch.long)
+        tokens = tokens.repeat(n_return_sequences, 1)
+
+        generated_tokens = tokens.to(self.device)
+        with torch.no_grad():
+            while generated_tokens.size(1) < length:
+
+                with torch.autocast(device_type=self.device, dtype=torch.bfloat16):
+                    if self.model_returns_logits:
+                        logits = self.model(generated_tokens)
+                    else:
+                        logits = self.model(generated_tokens).logits
+
+                logits = logits[:, -1, :]  # Get last token logits (B, vocab_size)
+                probs = F.softmax(logits, dim=-1)  # Convert to probabilities
+
+                topk_probs, topk_indices = torch.topk(probs, k=10, dim=-1)
+                sampled_indices = torch.multinomial(topk_probs, 1)  # Shape: (B, 1)
+                next_tokens = torch.gather(topk_indices, -1, sampled_indices)  # (B, 1)
+
+                generated_tokens = torch.cat((generated_tokens, next_tokens), dim=1)
+
+        continuations = []
+        for i in range(n_return_sequences):
+            tokens = generated_tokens[i, :length].tolist()
+            decoded = self.tokenizer.decode(tokens)
+            print(f"=== sample {i} ===\n{decoded}")
+            continuations.append(decoded)
+        return continuations