Create generation.py

generation.py

@@ -0,0 +1,51 @@
+import torch
+from torch.nn import functional as F
+from gpt_class import GPTConfig, GPT
+# Assuming tiktoken is correctly imported and functions as expected
+import tiktoken
+
+# Setup device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model
+state_dict = torch.load('model_51999.pt', map_location=device)
+config = state_dict['config']
+model = GPT(config)
+model.load_state_dict(state_dict['model'])
+model.to(device)
+model.eval()
+
+# Set seed for reproducibility
+torch.manual_seed(42)
+torch.cuda.manual_seed_all(42)
+
+# Get tokenizer
+tokenizer = tiktoken.get_encoding("gpt2")
+
+def Generate(model, tokenizer, example, num_return_sequences, max_length):
+    model.eval()
+    tokens = tokenizer.encode(example)
+    tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).repeat(num_return_sequences, 1)
+    tokens = tokens.to(device)
+    sample_rng = torch.Generator(device=device)
+
+    xgen = tokens
+    while xgen.size(1) < max_length:
+        with torch.no_grad():
+            with torch.autocast(device_type=device):
+                logits, _ = model(xgen)  # Assumes model returns logits and optional loss
+            logits = logits[:, -1, :]  # Get last token logits
+            probs = F.softmax(logits, dim=-1)
+            topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
+            ix = torch.multinomial(topk_probs, 1, generator=sample_rng)
+            xcol = torch.gather(topk_indices, -1, ix)
+            xgen = torch.cat((xgen, xcol), dim=1)
+
+    # Generate output for each sequence
+    for i in range(num_return_sequences):
+        tokens = xgen[i, :max_length].tolist()
+        decoded = tokenizer.decode(tokens)
+        print(f"Sample {i+1}: {decoded}")
+
+# Generate text
+Generate(model, tokenizer, example="It is raining outside and", num_return_sequences=4, max_length=64)