Upload llm_trainer.py

llm_trainer.py

@@ -0,0 +1,80 @@
+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